diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..486a232 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.zip filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index f42b50c..ae9b556 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,5 @@ __pycache__/ # IDE files .idea .vs_code/ + +data/ diff --git a/perception/__init__.py b/perception/__init__.py index 522e077..74b1032 100644 --- a/perception/__init__.py +++ b/perception/__init__.py @@ -1,8 +1,14 @@ -import perception.tasks.TestTasks.TestAlgo as TestAlgo -import perception.tasks.gate.GateCenter as GateSeg +import perception.vis.TestAlgo as TestAlgo +import perception.tasks.gate.GateCenterAlgo as GateSeg +import perception.tasks.gate.GateSegmentationAlgoA as GateSegA +import perception.tasks.gate.GateSegmentationAlgoB as GateSegB +import perception.tasks.gate.GateSegmentationAlgoC as GateSegC # import perception.tasks as tasks ALGOS = { 'test': TestAlgo.TestAlgo, - 'gateseg': GateSeg.GateCenter + 'gateseg': GateSeg.GateCenterAlgo, + 'gatesegA': GateSegA.GateSegmentationAlgoA, + 'gatesegB': GateSegB.GateSegmentationAlgoB, + 'gatesegC': GateSegC.GateSegmentationAlgoC } diff --git a/perception/tasks/TaskPerceiver.py b/perception/tasks/TaskPerceiver.py index 605c090..6c55e69 100644 --- a/perception/tasks/TaskPerceiver.py +++ b/perception/tasks/TaskPerceiver.py @@ -1,8 +1,8 @@ -from typing import Any, Dict, Tuple +from typing import Any, Dict import numpy as np -class TaskPerceiver: +class TaskPerceiver: def __init__(self, **kwargs): """Initializes the TaskPerceiver. Args: @@ -11,22 +11,18 @@ def __init__(self, **kwargs): for the slider which controls this variable, and default_val is the initial value of the slider. """ - self.time = 0 self.kwargs = kwargs def analyze(self, frame: np.ndarray, debug: bool, slider_vals: Dict[str, int]) -> Any: """Runs the algorithm and returns the result. Args: - frame: The frame to analyze + frame: The frame to analyze debug: Whether or not to display intermediate images for debugging slider_vals: A list of names of the variables which the user should be able to control from the Visualizer, mapped to current slider value for that variable - Returns: - the result of the algorithm - debug frames must each be same size as original input frame. Might change this in the future. + Returns: + the result of the algorithm + debug frames must each be same size as original input frame. Might change this in the future. """ raise NotImplementedError("Need to implement with child class.") - - def var_info(self) -> Dict[str, Tuple[Tuple[int, int], int]]: - return self.kwargs \ No newline at end of file diff --git a/perception/tasks/cross/CrossPerceiver.py b/perception/tasks/cross/CrossPerceiver.py deleted file mode 100644 index 5c12123..0000000 --- a/perception/tasks/cross/CrossPerceiver.py +++ /dev/null @@ -1,8 +0,0 @@ -from collections import namedtuple -import numpy as np -import sys -from perception.tasks.TaskPerceiver import TaskPerceiver - -class CrossPerceiver(TaskPerceiver): - named_tuple = namedtuple("CrossOutput", ["centerx", "centery"]) - named_tuple_types = {centerx: np.int16, centery: np.int16} diff --git a/perception/tasks/TestTasks/__init__.py b/perception/tasks/cross/__init__.py similarity index 100% rename from perception/tasks/TestTasks/__init__.py rename to perception/tasks/cross/__init__.py diff --git a/perception/tasks/cross/cross_detection.py b/perception/tasks/cross/cross_detection.py index 0b49772..9410f4f 100644 --- a/perception/tasks/cross/cross_detection.py +++ b/perception/tasks/cross/cross_detection.py @@ -8,10 +8,10 @@ ############################################################################# sys.path.insert(0, '../background_removal') -from peak_removal_adaptive_thresholding import filter_out_highest_peak_multidim -from combined_filter import combined_filter +from perception.tasks.segmentation.peak_removal_adaptive_thresholding import filter_out_highest_peak_multidim +from perception.tasks.segmentation.combinedFilter import init_combined_filter -ret, frame = True, cv2.imread('../data/cross/cross.png') # https://i.imgur.com/rjv1Vcy.png +ret, frame = True, cv2.imread('../data/cross/cross.png') # https://i.imgur.com/rjv1Vcy.png # "hsv" = Apply hsv thresholding before trying to find the path marker # "multidim" = Apply filter_out_highest_peak_multidim @@ -29,7 +29,7 @@ def find_cross(frame, draw_figs=True): gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) ret, thresh = cv2.threshold(gray, 127, 255,0) - __, contours,hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) + __, contours,hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) contours.sort(key=lambda c: cv2.contourArea(c), reverse=True) possible_crosses = [] @@ -44,17 +44,16 @@ def find_cross(frame, draw_figs=True): if defects is not None and len(defects) == 4: possible_crosses.append(defects) - if draw_figs: img = frame.copy() for defects in possible_crosses: for i in range(defects.shape[0]): - s,e,f,d = defects[i,0] + s, e, f, d = defects[i, 0] # start = tuple(cnt[s][0]) # end = tuple(cnt[e][0]) far = tuple(cnt[f][0]) # cv2.line(img,start,end,[0,255,0],2) - cv2.circle(img,far,5,[0,0,255],-1) + cv2.circle(img, far, 5, [0, 0, 255], -1) cv2.imshow('cross at contour number ' + str(i),img) cv2.imshow('original', frame) @@ -64,15 +63,15 @@ def find_cross(frame, draw_figs=True): ########################################### # Main Body ########################################### - +# TODO: port to vis if __name__ == "__main__": + combined_filter = init_combined_filter() + ret_tries = 0 - while(1 and ret_tries < 50): + while 1 and ret_tries < 50: # ret,frame = cap.read() - - if ret == True: + if ret: # frame = cv2.resize(frame, (0,0), fx=0.5, fy=0.5) - if thresholding == "multidim": votes1, threshed = filter_out_highest_peak_multidim(frame) threshed = cv2.morphologyEx(threshed, cv2.MORPH_OPEN, np.ones((5,5),np.uint8)) @@ -86,13 +85,9 @@ def find_cross(frame, draw_figs=True): ret_tries = 0 k = cv2.waitKey(60) & 0xff - if k == 27: # esc - if testing: - print("hsv thresholds:") - print(thresholds_used) + if k == 27: # esc break else: ret_tries += 1 cv2.destroyAllWindows() - cap.release() diff --git a/perception/tasks/gate/GateCenter.py b/perception/tasks/gate/GateCenterAlgo.py similarity index 57% rename from perception/tasks/gate/GateCenter.py rename to perception/tasks/gate/GateCenterAlgo.py index 4f90859..5c3e6f6 100644 --- a/perception/tasks/gate/GateCenter.py +++ b/perception/tasks/gate/GateCenterAlgo.py @@ -1,16 +1,14 @@ -from perception.tasks.gate.GateSegmentation import GateSegmentationAlgo +from perception.tasks.gate.GateSegmentationAlgoA import GateSegmentationAlgoA from perception.tasks.TaskPerceiver import TaskPerceiver from collections import namedtuple -import sys import numpy as np import math import cv2 as cv -import time import statistics -class GateCenter(TaskPerceiver): +class GateCenterAlgo(TaskPerceiver): center_x_locs, center_y_locs = [], [] output_class = namedtuple("GateOutput", ["centerx", "centery"]) output_type = {'centerx': np.int16, 'centery': np.int16} @@ -20,39 +18,43 @@ def __init__(self): self.gate_center = self.output_class(250, 250) self.use_optical_flow = False self.optical_flow_c = 0.1 - self.gate = GateSegmentationAlgo() + self.gate = GateSegmentationAlgoA() self.prvs = None + # TODO: do input and return typing def analyze(self, frame, debug, slider_vals): - self.optical_flow_c = slider_vals['optical_flow_c'] / 100 - rect1, rect2, debug_filter = self.gate.analyze(frame, True) + self.optical_flow_c = slider_vals['optical_flow_c']/100 + rect, debug_filters = self.gate.analyze(frame, True) + debug_filter = debug_filters[-1] + debug_filters = debug_filters[:-1] + if self.prvs is None: # frame = cv.resize(frame, None, fx=0.3, fy=0.3) self.prvs = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) - else: - if rect1 and rect2: - self.gate_center = self.get_center(rect1, rect2, frame) + else: + if rect[0] and rect[1]: + self.gate_center = self.get_center(rect[0], rect[1], frame) if self.use_optical_flow: - cv.circle(debug_filter, self.gate_center, 5, (3, 186, 252), -1) + cv.circle(debug_filter, self.gate_center, 5, (3,186,252), -1) else: - cv.circle(debug_filter, self.gate_center, 5, (0, 0, 255), -1) - + cv.circle(debug_filter, self.gate_center, 5, (0,0,255), -1) + if debug: - return (self.output_class(self.gate_center[0], self.gate_center[1]), [frame, debug_filter]) - return self.output_class(self.gate_center[0], self.gate_center[1]) + return (self.gate_center[0], self.gate_center[1]), list(debug_filters) + [debug_filter] + return (self.gate_center[0], self.gate_center[1]) def center_without_optical_flow(self, center_x, center_y): # get starting center location, averaging over the first 2510 frames if len(self.center_x_locs) == 0: self.center_x_locs.append(center_x) self.center_y_locs.append(center_y) - + elif len(self.center_x_locs) < 25: self.center_x_locs.append(center_x) self.center_y_locs.append(center_y) center_x = int(statistics.mean(self.center_x_locs)) center_y = int(statistics.mean(self.center_y_locs)) - + # use new center location only when it is close to the previous valid location else: self.center_x_locs.append(center_x) @@ -61,11 +63,11 @@ def center_without_optical_flow(self, center_x, center_y): self.center_y_locs.pop(0) x_temp_avg = int(statistics.mean(self.center_x_locs)) y_temp_avg = int(statistics.mean(self.center_y_locs)) - if math.sqrt((center_x - x_temp_avg) ** 2 + (center_y - y_temp_avg) ** 2) > 10: + if math.sqrt((center_x - x_temp_avg)**2 + (center_y - y_temp_avg)**2) > 10: center_x, center_y = int(x_temp_avg), int(y_temp_avg) - + return (center_x, center_y) - + def dense_optical_flow(self, frame): next_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) flow = cv.calcOpticalFlowFarneback(self.prvs, next_frame, None, 0.5, 3, 15, 3, 5, 1.2, 0) @@ -82,12 +84,9 @@ def get_center(self, rect1, rect2, frame): x2, y2, w2, h2 = rect2 center_x, center_y = (x1 + x2) // 2, ((y1 + h1 // 2) + (y2 + h2 // 2)) // 2 self.prvs, mag, ang = self.dense_optical_flow(frame) - # print(np.mean(mag)) - if len(self.center_x_locs) < 25 or (np.mean(mag) < 40 and ((not self.use_optical_flow) or \ - (self.use_optical_flow and ( - center_x - self.gate_center[0]) ** 2 + ( - center_y - self.gate_center[ - 1]) ** 2 < 50))): + + if len(self.center_x_locs) < 25 or (np.mean(mag) < 40 and ((not self.use_optical_flow ) or \ + (self.use_optical_flow and (center_x - self.gate_center[0])**2 + (center_y - self.gate_center[1])**2 < 50))): self.use_optical_flow = False return self.center_without_optical_flow(center_x, center_y) self.use_optical_flow = True @@ -95,40 +94,6 @@ def get_center(self, rect1, rect2, frame): (int(self.gate_center[1] + self.optical_flow_c * np.mean(mag * np.sin(ang))))) -# this part is temporary and will be covered by other files in the future if __name__ == '__main__': - cap = cv.VideoCapture(sys.argv[1]) - ret_tries = 0 - start_time = time.time() - frame_count = 0 - paused = False - speed = 1 - ret, frame1 = cap.read() - frame1 = cv.resize(frame1, None, fx=0.3, fy=0.3) - prvs = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY) - hsv = np.zeros_like(frame1) - hsv[..., 1] = 255 - gate_center = GateCenter() - while ret_tries < 50: - for _ in range(speed): - ret, frame = cap.read() - if frame_count == 1000: - break - if ret: - frame = cv.resize(frame, None, fx=0.3, fy=0.3) - center, filtered_frame = gate_center.analyze(frame, True) - cv.imshow('original', frame) - cv.imshow('filtered_frame', filtered_frame) - ret_tries = 0 - key = cv.waitKey(30) - if key == ord('q') or key == 27: - break - if key == ord('p'): - paused = not paused - if key == ord('i') and speed > 1: - speed -= 1 - if key == ord('o'): - speed += 1 - else: - ret_tries += 1 - frame_count += 1 \ No newline at end of file + from perception.vis.vis import run + run(['..\..\..\data\GOPR1142.MP4'], GateCenterAlgo(), False) \ No newline at end of file diff --git a/perception/tasks/gate/GateSegmentation.py b/perception/tasks/gate/GateSegmentationAlgoA.py similarity index 62% rename from perception/tasks/gate/GateSegmentation.py rename to perception/tasks/gate/GateSegmentationAlgoA.py index 2c4c052..5e3173f 100644 --- a/perception/tasks/gate/GateSegmentation.py +++ b/perception/tasks/gate/GateSegmentationAlgoA.py @@ -1,29 +1,28 @@ from perception.tasks.TaskPerceiver import TaskPerceiver from typing import Tuple -import sys from perception.tasks.segmentation.combinedFilter import init_combined_filter import numpy as np import cv2 as cv -import time -class GateSegmentationAlgo(TaskPerceiver): +class GateSegmentationAlgoA(TaskPerceiver): center_x_locs, center_y_locs = [], [] - + def __init__(self): super().__init__() self.combined_filter = init_combined_filter() + # TODO: fix return typing def analyze(self, frame: np.ndarray, debug: bool, slider_vals=None) -> Tuple[float, float]: """Takes in the background removed image and returns the center between the two gate posts. Args: frame: The background removed frame to analyze debug: Whether or not tot display intermediate images for debugging - Reurns: - (x,y) coordinate with center of gate - """ + Reurns: + (x,y) coordinate with center of gate + """ rect1, rect2 = None, None filtered_frame = self.combined_filter(frame, display_figs=False) @@ -33,19 +32,18 @@ def analyze(self, frame: np.ndarray, debug: bool, slider_vals=None) -> Tuple[flo upperbound = 255 _, thresh = cv.threshold(filtered_frame, lowerbound, upperbound, cv.THRESH_BINARY) debug_filter = cv.cvtColor(thresh, cv.COLOR_GRAY2BGR) - + cnt = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[-2] - + area_diff = [] area_cnts = [] # remove all contours with zero area cnt = [cnt[i] for i in range(len(cnt)) if cv.contourArea(cnt[i]) > 0] - - for i in cnt: - area_cnt = cv.contourArea(i) + for c in cnt: + area_cnt = cv.contourArea(c) area_cnts.append(area_cnt) - area_rect = cv.boundingRect(i)[-2] * cv.boundingRect(i)[-1] + area_rect = cv.boundingRect(c)[-2] * cv.boundingRect(c)[-1] area_diff.append(abs((area_rect - area_cnt) / area_cnt)) if len(area_diff) >= 2: @@ -61,39 +59,10 @@ def analyze(self, frame: np.ndarray, debug: bool, slider_vals=None) -> Tuple[flo cv.rectangle(debug_filter, (x2, y2), (x2 + w2, y2 + h2), (0, 255, 0), 2) if debug: - return (rect1, rect2, debug_filter) + return (rect1, rect2), (frame, debug_filter) return (rect1, rect2) + - -# this part is temporary and will be covered by other files in the future if __name__ == '__main__': - cap = cv.VideoCapture(sys.argv[1]) - ret_tries = 0 - start_time = time.time() - frame_count = 0 - paused = False - speed = 1 - gate_task = GateSegmentationAlgo() - while ret_tries < 50: - for _ in range(speed): - ret, frame = cap.read() - if frame_count == 1000: - break - if ret: - frame = cv.resize(frame, None, fx=0.3, fy=0.3) - rect1, rect2, filtered_frame = gate_task.analyze(frame, True) - cv.imshow('original', frame) - cv.imshow('filtered_frame', filtered_frame) - ret_tries = 0 - key = cv.waitKey(30) - if key == ord('q') or key == 27: - break - if key == ord('p'): - paused = not paused - if key == ord('i') and speed > 1: - speed -= 1 - if key == ord('o'): - speed += 1 - else: - ret_tries += 1 - frame_count += 1 \ No newline at end of file + from perception.vis.vis import run + run(['..\..\..\data\GOPR1142.MP4'], GateSegmentationAlgoA(), False) diff --git a/perception/tasks/gate/GateSegmentationAlgoB.py b/perception/tasks/gate/GateSegmentationAlgoB.py new file mode 100644 index 0000000..8f0a54e --- /dev/null +++ b/perception/tasks/gate/GateSegmentationAlgoB.py @@ -0,0 +1,94 @@ +from typing import Tuple + +from perception.tasks.TaskPerceiver import TaskPerceiver +from perception.tasks.segmentation.combinedFilter import init_combined_filter +import numpy as np +import cv2 as cv +import statistics + + +class GateSegmentationAlgoB(TaskPerceiver): + center_x_locs, center_y_locs = [], [] + + def __init__(self): + super().__init__() + self.combined_filter = init_combined_filter() + + def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: + """Takes in the background removed image and returns the center between + the two gate posts. + Args: + frame: The background removed frame to analyze + debug: Whether or not tot display intermediate images for debugging + Reurns: + (x,y) coordinate with center of gate + """ + gate_center = (250, 250) + filtered_frame = self.combined_filter(frame, display_figs=False) + + max_brightness = max([b for b in filtered_frame[:, :, 0][0]]) + lowerbound = max(0.84*max_brightness, 120) + upperbound = 255 + _,thresh = cv.threshold(filtered_frame,lowerbound, upperbound, cv.THRESH_BINARY) + debug_filter = cv.cvtColor(thresh, cv.COLOR_GRAY2BGR) + + cnt = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[-2] + + area_diff = [] + area_cnts = [] + + # remove all contours with zero area + cnt = [cnt[i] for i in range(len(cnt)) if cv.contourArea(cnt[i]) > 0] + + for c in cnt: + area_cnt = cv.contourArea(c) + area_cnts.append(area_cnt) + area_rect = cv.boundingRect(c)[-2] * cv.boundingRect(c)[-1] + area_diff.append(abs((area_rect - area_cnt)/area_cnt)) + + if len(area_diff) >= 2: + largest_area_idx = [area_cnts.index(sorted(area_cnts, reverse=True)[i]) for i in range(min(3, len(cnt)))] + area_diff_copy = sorted([area_diff[i] for i in largest_area_idx]) + min_i1, min_i2 = area_diff.index(area_diff_copy[0]), area_diff.index(area_diff_copy[1]) + + (x1, y1, w1, h1) = cv.boundingRect(cnt[min_i1]) + (x2, y2, w2, h2) = cv.boundingRect(cnt[min_i2]) + cv.rectangle(debug_filter, (x1, y1), (x1+w1, y1+h1), (0,255,0), 2) + cv.rectangle(debug_filter, (x2, y2), (x2+w2, y2+h2), (0,255,0), 2) + + # drawing center dot + center_x, center_y = (x1+x2)//2, ((y1+h1//2)+(y2+h2//2))//2 + gate_center = self.get_actual_center(center_x, center_y) + cv.circle(debug_filter, gate_center, 5, (0,0,255), -1) + + if debug: + return (gate_center[0], gate_center[1]), (frame, debug_filter) + return (gate_center[0], gate_center[1]) + + def get_actual_center(self, center_x, center_y): + # get starting center location, averaging over the first 2510 frames + if len(self.center_x_locs) == 0: + self.center_x_locs.append(center_x) + self.center_y_locs.append(center_y) + + elif len(self.center_x_locs) < 25: + self.center_x_locs.append(center_x) + self.center_y_locs.append(center_y) + center_x = int(statistics.mean(self.center_x_locs)) + center_y = int(statistics.mean(self.center_y_locs)) + + # use new center location only when it is close to the previous valid location + else: + if abs(center_x - self.center_x_locs[-1]) > 10 or \ + abs(center_y - self.center_y_locs[-1]) > 10: + center_x, center_y = self.center_x_locs[-1], self.center_y_locs[-1] + else: + self.center_x_locs.append(center_x) + self.center_y_locs.append(center_y) + + return (center_x, center_y) + + +if __name__ == '__main__': + from perception.vis.vis import run + run(['..\..\..\data\GOPR1142.MP4'], GateSegmentationAlgoB(), False) diff --git a/perception/tasks/TestTasks/GateSegmentationAlgo.py b/perception/tasks/gate/GateSegmentationAlgoC.py similarity index 55% rename from perception/tasks/TestTasks/GateSegmentationAlgo.py rename to perception/tasks/gate/GateSegmentationAlgoC.py index f747dd2..f3bdc55 100644 --- a/perception/tasks/TestTasks/GateSegmentationAlgo.py +++ b/perception/tasks/gate/GateSegmentationAlgoC.py @@ -1,35 +1,33 @@ from perception.tasks.TaskPerceiver import TaskPerceiver from typing import Tuple -import sys -import os -from pathlib import Path from collections import namedtuple from perception.tasks.segmentation.combinedFilter import init_combined_filter import numpy as np import cv2 as cv -import time -class GateSegmentationAlgo(TaskPerceiver): + +class GateSegmentationAlgoC(TaskPerceiver): __past_centers = [] __ema = None output_class = namedtuple("GateOutput", ["centerx", "centery"]) output_type = {'centerx': np.int16, 'centery': np.int16} def __init__(self, alpha=0.1): - super() + super().__init__() self.__alpha = alpha self.combined_filter = init_combined_filter() + # TODO: fix return typing def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: """Takes in the background removed image and returns the center between - the two gate posts. - Args: - frame: The background removed frame to analyze. - debug: Whether or not to display intermediate images for debugging. - Returns: - (x,y) coordinate with center of gate - """ + the two gate posts. + Args: + frame: The background removed frame to analyze. + debug: Whether or not to display intermediate images for debugging. + Returns: + (x,y) coordinate with center of gate + """ gate_center = self.output_class(250, 250) filtered_frame = self.combined_filter(frame, display_figs=False) filtered_frame_copies = [filtered_frame for _ in range(3)] @@ -37,9 +35,9 @@ def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: mask = cv.inRange( stacked_filter_frames, np.array([100, 100, 100]), np.array([255, 255, 255]) ) - _, contours, _ = cv.findContours(mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) - if contours: - contours.sort(key=self.findStraightness, reverse=True) + contours = cv.findContours(mask, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)[-2] + if len(contours) > 0: + contours.sort(reverse=True, key=self.findStraightness) cnts = contours[:2] rects = [cv.minAreaRect(c) for c in cnts] centers = [np.array(r[0]) for r in rects] @@ -56,6 +54,7 @@ def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: self.__alpha * gate_center + (1 - self.__alpha) * self.__ema ) gate_center = (int(self.__ema[0]), int(self.__ema[1])) + # TODO: clean this up via hyperparam or move to gate center algo # if len(self.__past_centers) < 15: # self.__past_centers += [gate_center] # else: @@ -66,53 +65,16 @@ def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: cv.circle(stacked_filter_frames, gate_center, 10, (0, 255, 0), -1) if debug: - return ( - self.output_class(gate_center[0], gate_center[1]), - [stacked_filter_frames], - ) - return self.output_class(gate_center[0], gate_center[1]) + return (gate_center[0], gate_center[1]), (frame, stacked_filter_frames) + return (gate_center[0], gate_center[1]) - def findStraightness( - self, contour - ): # output number = contour area/convex area, the bigger the straightest + def findStraightness(self, contour): # output number = contour area/convex area, the bigger the straightest hull = cv.convexHull(contour, False) contour_area = cv.contourArea(contour) hull_area = cv.contourArea(hull) return 10 * contour_area - 5 * hull_area -# this part is temporary and will be covered by other files in the future if __name__ == '__main__': - combined_filter = init_combined_filter() - cap = cv.VideoCapture(sys.argv[1]) - ret_tries = 0 - gate_task = GateSegmentationAlgo(0.1) - # once = False - start_time = time.time() - frame_count = 0 - while ret_tries < 50: - ret, frame = cap.read() - if frame_count == 1000: - break - if ret: - frame = cv.resize(frame, None, fx=0.4, fy=0.4) - - ### FUNCTION CALL, can change this - center, filtered_frame = gate_task.analyze(frame, True) - # cProfile.run("gate_task.analyze(frame, True)") - # cv.putText(frame, "x: %.2f" % x + " y: %.2f" % y, - # (20, frame.shape[0] - 20), cv.FONT_HERSHEY_SIMPLEX, - # 2.0, (0, 165, 255), 3) - cv.imshow('original', frame) - cv.imshow('filtered_frame', filtered_frame) - # if not once: - # print(filtered_frame) - # once = True - ret_tries = 0 - k = cv.waitKey(60) & 0xFF - if k == 27: - break - else: - ret_tries += 1 - frame_count += 1 - # print(frame_count / (time.time() - start_time)) + from perception.vis.vis import run + run(['..\..\..\data\GOPR1142.MP4'], GateSegmentationAlgoC(), False) diff --git a/perception/tasks/gate/archive/detectGate.py b/perception/tasks/gate/archive/detectGate.py index cf309ff..5f7e868 100644 --- a/perception/tasks/gate/archive/detectGate.py +++ b/perception/tasks/gate/archive/detectGate.py @@ -1,10 +1,10 @@ import numpy as np import cv2 import argparse -import sys -from PIL import Image import time +# TODO: port to vis + TaskPerciever format or remove + video_file = 'truncated_semi_final_run.mp4' EPSILON = 40 OVERLAP_EPS = 40 diff --git a/perception/tasks/gate/archive/threshTest.py b/perception/tasks/gate/archive/threshTest.py index 0bd6c71..70d8aa5 100644 --- a/perception/tasks/gate/archive/threshTest.py +++ b/perception/tasks/gate/archive/threshTest.py @@ -2,6 +2,9 @@ import cv2 as cv import argparse import numpy as np + +# TODO: port to vis + TaskPerciever format or remove + #expectations #contours closest to the last ones #should know when we passed through the gate diff --git a/perception/tasks/gate/legacy/GatePerceiver.py b/perception/tasks/gate/legacy/GatePerceiver.py deleted file mode 100644 index 566b48d..0000000 --- a/perception/tasks/gate/legacy/GatePerceiver.py +++ /dev/null @@ -1,9 +0,0 @@ -from collections import namedtuple -import numpy as np -import sys -sys.path.insert(0, '../..') -from TaskPerceiver import TaskPerceiver - -class GatePerceiver(TaskPerceiver): - output_class = namedtuple("GateOutput", ["centerx", "centery"]) - output_type = {'centerx': np.int16, 'centery': np.int16} diff --git a/perception/tasks/gate/legacy/GateSegmentationAlgo.py b/perception/tasks/gate/legacy/GateSegmentationAlgo.py deleted file mode 100644 index 191865c..0000000 --- a/perception/tasks/gate/legacy/GateSegmentationAlgo.py +++ /dev/null @@ -1,108 +0,0 @@ -from GatePerceiver import GatePerceiver -from typing import Tuple -import sys -import os -sys.path.append(os.path.dirname(__file__)) - - -from segmentation.combinedFilter import init_combined_filter -import numpy as np -import cv2 as cv -import time -import cProfile - -class GateSegmentationAlgo(GatePerceiver): - __past_centers = [] - __ema = None - - def __init__(self, alpha): - super() - self.__alpha = alpha - - def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: - """Takes in the background removed image and returns the center between - the two gate posts. - Args: - frame: The background removed frame to analyze - debug: Whether or not tot display intermediate images for debugging - Reurns: - (x,y) coordinate with center of gate - """ - gate_center = self.output_class(250, 250) - filtered_frame = combined_filter(frame, display_figs=False) - filtered_frame_copies = [filtered_frame for _ in range(3)] - stacked_filter_frames = np.concatenate(filtered_frame_copies, axis = 2) - mask = cv.inRange(stacked_filter_frames, - np.array([100, 100, 100]), np.array([255, 255, 255])) - _, contours, _ = cv.findContours(mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) - if contours: - contours.sort(key=self.findStraightness, reverse=True) - cnts = contours[:2] - rects = [cv.minAreaRect(c) for c in cnts] - centers = [np.array(r[0]) for r in rects] - boxpts = [cv.boxPoints(r) for r in rects] - box = [np.int0(b) for b in boxpts] - for b in box: - cv.drawContours(stacked_filter_frames,[b],0,(0,0,255),5) - if len(centers) >= 2: - gate_center = (centers[0] + centers[1]) * 0.5 - if self.__ema is None: - self.__ema = gate_center - else: - self.__ema = self.__alpha*gate_center + (1 - self.__alpha)*self.__ema - gate_center = (int(self.__ema[0]), int(self.__ema[1])) - # if len(self.__past_centers) < 15: - # self.__past_centers += [gate_center] - # else: - # self.__past_centers.pop(0) - # self.__past_centers += [gate_center] - # gate_center = sum(self.__past_centers) / len(self.__past_centers) - # gate_center = (int(gate_center[0]), int(gate_center[1])) - cv.circle(stacked_filter_frames, gate_center, 10, (0,255,0), -1) - - if debug: - return (self.output_class(gate_center[0], gate_center[1]), stacked_filter_frames) - return self.output_class(gate_center[0], gate_center[1]) - - def findStraightness(self, contour): # output number = contour area/convex area, the bigger the straightest - hull = cv.convexHull(contour, False) - contour_area = cv.contourArea(contour) - hull_area = cv.contourArea(hull) - return 10 * contour_area - 5 * hull_area - -# this part is temporary and will be covered by other files in the future -if __name__ == '__main__': - combined_filter = init_combined_filter() - cap = cv.VideoCapture(sys.argv[1]) - ret_tries = 0 - gate_task = GateSegmentationAlgo(0.1) - # once = False - start_time = time.time() - frame_count = 0 - while ret_tries < 50: - ret, frame = cap.read() - if frame_count == 1000: - break - if ret: - frame = cv.resize(frame, None, fx=0.4, fy=0.4) - - - ### FUNCTION CALL, can change this - center, filtered_frame = gate_task.analyze(frame, True) - # cProfile.run("gate_task.analyze(frame, True)") - # cv.putText(frame, "x: %.2f" % x + " y: %.2f" % y, - # (20, frame.shape[0] - 20), cv.FONT_HERSHEY_SIMPLEX, - # 2.0, (0, 165, 255), 3) - cv.imshow('original', frame) - cv.imshow('filtered_frame', filtered_frame) - # if not once: - # print(filtered_frame) - # once = True - ret_tries = 0 - k = cv.waitKey(60) & 0xff - if k == 27: - break - else: - ret_tries += 1 - frame_count += 1 - #print(frame_count / (time.time() - start_time)) diff --git a/perception/tasks/path_marker/PathMarkerPerceiver.py b/perception/tasks/path_marker/PathMarkerPerceiver.py deleted file mode 100644 index d77335d..0000000 --- a/perception/tasks/path_marker/PathMarkerPerceiver.py +++ /dev/null @@ -1,9 +0,0 @@ -from collections import namedtuple -import numpy as np -import sys -sys.path.insert(0, '..') -from TaskPerceiver import TaskPerceiver - -class PathMarkerPerceiver(TaskPerceiver): - named_tuple = namedtuple("PathMarkerOutput", ["angle"]) - named_tuple_types = {angle: np.float64} diff --git a/perception/tasks/gate/legacy/__init__.py b/perception/tasks/path_marker/__init__.py similarity index 100% rename from perception/tasks/gate/legacy/__init__.py rename to perception/tasks/path_marker/__init__.py diff --git a/perception/tasks/path_marker/path_marker_detection.py b/perception/tasks/path_marker/path_marker_detection.py index f788bd5..83eef2c 100644 --- a/perception/tasks/path_marker/path_marker_detection.py +++ b/perception/tasks/path_marker/path_marker_detection.py @@ -1,6 +1,8 @@ -from combined_filter import init_combined_filter +from perception.tasks.segmentation.combinedFilter import init_combined_filter from typing import Union +# TODO: port to vis + TaskPerciever format or remove + if __name__ == "__main__": import numpy as np import cv2 @@ -31,7 +33,7 @@ def thresh_by_contour_size( frame, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE ) if contours is not None: - contours.sort(key=lambda c: cv2.contourArea(c), reverse=True) + contours.sort(key=cv2.contourArea, reverse=True) contours = contours[:num_contours] threshed = np.zeros(frame.shape, np.uint8) @@ -126,13 +128,11 @@ def line_length(line): ) return bot_angle, top_angle - else: - if draw_figs: - cv2.imshow('lines bottom', frame) - cv2.imshow('lines top', frame) - cv2.imshow('frame with path marker angles', frame) - return None + if draw_figs: + cv2.imshow('lines bottom', frame) + cv2.imshow('lines top', frame) + cv2.imshow('frame with path marker angles', frame) def path_marker_get_new_heading(cap, is_approaching, draw_figs=False): diff --git a/perception/tasks/roulette/__init__.py b/perception/tasks/roulette/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/perception/tasks/roulette/spinny_wheel_detection.py b/perception/tasks/roulette/spinny_wheel_detection.py new file mode 100644 index 0000000..463c73a --- /dev/null +++ b/perception/tasks/roulette/spinny_wheel_detection.py @@ -0,0 +1,119 @@ +import numpy as np +import cv2 +import argparse +import sys +import time + +# TODO: port to vis + TaskPerciever format or remove + +# CHANGE PARAMETER IN CALL TO "THRESH" FUNCTION TO CHANGE COLOR +file_name = "GOPR1145.MP4" # video file from dropbox +vid = cv2.VideoCapture(file_name) +frames = 0 +avgLength = 10 +centers = [] + + +def thresh(frame, color='red'): + blur = cv2.GaussianBlur(frame, (5, 5), 0) + hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV) + + if color == 'red': + lower = np.uint8([29, 77, 36]) + upper = np.uint8([130, 250, 255]) + mask = cv2.inRange(hsv, lower, upper) + mask = cv2.bitwise_not(mask) + elif color == 'blue': + lower = np.uint8([86, 141, 0]) + upper = np.uint8([106, 220, 168]) + mask = cv2.inRange(hsv, lower, upper) + else: + lower = np.uint8([66, 208, 157]) + upper = np.uint8([86, 255, 209]) + mask = cv2.inRange(hsv, lower, upper) + + return mask + + +def heuristic(contour): + rect = cv2.minAreaRect(contour) + area = rect[1][0] * rect[1][1] + diff = cv2.contourArea(cv2.convexHull(contour)) - cv2.contourArea(contour) + cent = rect[0] + dist = 0 + if len(likelySection) > 1 and allLarger(60): + cen0 = cv2.minAreaRect(likelySection[0]['cont'])[0] + dis0 = np.linalg.norm(np.array(cent) - np.array(cen0)) + cen1 = cv2.minAreaRect(likelySection[1]['cont'])[0] + dis1 = np.linalg.norm(np.array(cent) - np.array(cen1)) + dist = min([dis0, dis1]) + heur = area - 3 * diff - 20 * dist + return heur + + +def allLarger(thresh): + for cnt in likelySection: + if cnt['heur'] < thresh: + return False + return True + + +def drawRects(frame, contours): + tempPts = [] + for cnt in contours: + rect = cv2.minAreaRect(cnt['cont']) + boxpts = cv2.boxPoints(rect) + box = np.int0(boxpts) + cv2.drawContours(frame, [box], 0, (0, 0, 255), 1) + cv2.drawContours(frame, [cnt['cont']], 0, (0, 255, 0), 1) + cv2.drawContours(frame, [cv2.convexHull(cnt['cont'])], 0, (255, 0, 0), 1) + tempPts.append(rect[0]) + # cv2.putText(frame, str(cnt['heur']), (int(rect[0][0]), int(rect[0][1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255)) + if len(tempPts) > 1 and allLarger(60): + cv2.circle(frame, (int(tempPts[0][0]), int(tempPts[0][1])), 10, (0, 0, 255), -1) + cv2.circle(frame, (int(tempPts[1][0]), int(tempPts[1][1])), 10, (0, 0, 255), -1) + + +def midPt(pt1, pt2): + return ((pt1[0] + pt2[0]) / 2, (pt1[1] + pt2[1]) / 2) + + +def getAvgPt(pt): + points.append(pt) + exes = list(map(lambda x: x[0], points)) + whys = list(map(lambda y: y[1], points)) + + if len(points) > 50: + del points[:10] + return (int(sum(exes) / len(exes)), int(sum(whys) / len(whys))) + + +likelySection = [] +points = [] +while vid.isOpened(): + start = time.time() + ret, frame = vid.read() + if (ret == False): + continue + + threshed = thresh(frame, 'other') + res, contours, hierarchy = cv2.findContours(threshed, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + + contours.sort(key=heuristic, reverse=True) + if len(contours) > 1: + c1 = contours[0] + c2 = contours[1] + heur0 = heuristic(c1) + heur1 = heuristic(c2) + likelySection = [{'cont': contours[0], 'heur': heur0}, {'cont': contours[1], 'heur': heur1}] + untampered = np.copy(frame) + if contours: + drawRects(frame, likelySection) + cv2.imshow("Frame", frame) + cv2.imshow('Res', res) + + if (cv2.waitKey(1) & 0xFF) == ord('q') or frames > 900: + break + +vid.release() +cv2.destroyAllWindows() diff --git a/perception/tasks/spinny/threshslider.py b/perception/tasks/roulette/threshslider.py similarity index 66% rename from perception/tasks/spinny/threshslider.py rename to perception/tasks/roulette/threshslider.py index e8bf654..fadc915 100644 --- a/perception/tasks/spinny/threshslider.py +++ b/perception/tasks/roulette/threshslider.py @@ -2,20 +2,23 @@ import cv2 as cv import argparse import numpy as np -#expectations -#contours closest to the last ones -#should know when we passed through the gate + +# TODO: port to vis + TaskPerciever format or remove + +# expectations +# contours closest to the last ones +# should know when we passed through the gate """ IMPORTANT!!!! RUN THIS WITH $ python3 threshTest.py GOPR1142.mp4 """ max_value = 255 -max_value_H = 360//2 -low_H = 86#49#29#0 -low_S = 141#77#0 -low_V = 0#36#0 For Small sector, increasing lower V bound reduces -high_H = 106#130#max_value_H -high_S = 217#250#max_value -high_V = 168#max_value +max_value_H = 360 // 2 +low_H = 86 # 49#29#0 +low_S = 141 # 77#0 +low_V = 0 # 36#0 For Small sector, increasing lower V bound reduces +high_H = 106 # 130#max_value_H +high_S = 217 # 250#max_value +high_V = 168 # max_value window_capture_name = 'Video Capture' window_detection_name = 'Object Detection' low_H_name = 'Low H' @@ -24,63 +27,78 @@ high_H_name = 'High H' high_S_name = 'High S' high_V_name = 'High V' + + def on_low_H_thresh_trackbar(val): global low_H global high_H low_H = val - low_H = min(high_H-1, low_H) + low_H = min(high_H - 1, low_H) cv.setTrackbarPos(low_H_name, window_detection_name, low_H) + + def on_high_H_thresh_trackbar(val): global low_H global high_H high_H = val - high_H = max(high_H, low_H+1) + high_H = max(high_H, low_H + 1) cv.setTrackbarPos(high_H_name, window_detection_name, high_H) + + def on_low_S_thresh_trackbar(val): global low_S global high_S low_S = val - low_S = min(high_S-1, low_S) + low_S = min(high_S - 1, low_S) cv.setTrackbarPos(low_S_name, window_detection_name, low_S) + + def on_high_S_thresh_trackbar(val): global low_S global high_S high_S = val - high_S = max(high_S, low_S+1) + high_S = max(high_S, low_S + 1) cv.setTrackbarPos(high_S_name, window_detection_name, high_S) + + def on_low_V_thresh_trackbar(val): global low_V global high_V low_V = val - low_V = min(high_V-1, low_V) + low_V = min(high_V - 1, low_V) cv.setTrackbarPos(low_V_name, window_detection_name, low_V) + + def on_high_V_thresh_trackbar(val): global low_V global high_V high_V = val - high_V = max(high_V, low_V+1) + high_V = max(high_V, low_V + 1) cv.setTrackbarPos(high_V_name, window_detection_name, high_V) + def drawRects(frame, contours): tempPts = [] for cnt in contours: rect = cv.minAreaRect(cnt['cont']) boxpts = cv.boxPoints(rect) box = np.int0(boxpts) - cv.drawContours(frame,[box],0,(0,0,255),1) - cv.drawContours(frame, [cnt['cont']],0,(0,255,0),1) - cv.drawContours(frame, [cv.convexHull(cnt['cont'])],0,(255,0,0),1) + cv.drawContours(frame, [box], 0, (0, 0, 255), 1) + cv.drawContours(frame, [cnt['cont']], 0, (0, 255, 0), 1) + cv.drawContours(frame, [cv.convexHull(cnt['cont'])], 0, (255, 0, 0), 1) tempPts.append(rect[0]) cv.putText(frame, str(cnt['heur']), (int(rect[0][0]), int(rect[0][1])), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255)) if len(tempPts) > 1 and allLarger(60): - #global paused + # global paused paused = True avgPt = getAvgPt(midPt(tempPts[0], tempPts[1])) - cv.circle(frame, (avgPt[0], avgPt[1]), 10, (0,0,255), -1) + cv.circle(frame, (avgPt[0], avgPt[1]), 10, (0, 0, 255), -1) + def midPt(pt1, pt2): return ((pt1[0] + pt2[0]) / 2, (pt1[1] + pt2[1]) / 2) + def getAvgPt(pt): points.append(pt) exes = list(map(lambda x: x[0], points)) @@ -90,6 +108,7 @@ def getAvgPt(pt): del points[:10] return (int(sum(exes) / len(exes)), int(sum(whys) / len(whys))) + """ def findLikelyGate(rectList, contours): if not (rectList and contours): @@ -101,6 +120,7 @@ def findLikelyGate(rectList, contours): return """ + def heuristic(contour): rect = cv.minAreaRect(contour) area = rect[1][0] * rect[1][1] @@ -113,16 +133,18 @@ def heuristic(contour): cen1 = cv.minAreaRect(likelyGate[1]['cont'])[0] dis1 = np.linalg.norm(np.array(cent) - np.array(cen1)) dist = min([dis0, dis1]) - heur = area - 3 * diff - 20 * dist #only factor in dist with all heurs larger than 60 - #print(heur) + heur = area - 3 * diff - 20 * dist # only factor in dist with all heurs larger than 60 + # print(heur) return heur + def allLarger(thresh): for cnt in likelyGate: if cnt['heur'] < thresh: return False return True + parser = argparse.ArgumentParser(description='Code for Thresholding Operations using inRange tutorial.') parser.add_argument('camera', help='Camera devide number.', default=0, type=str) args = parser.parse_args() @@ -131,14 +153,14 @@ def allLarger(thresh): cv.namedWindow(window_capture_name) cv.namedWindow(window_detection_name) -cv.createTrackbar(low_H_name, window_detection_name , low_H, max_value_H, on_low_H_thresh_trackbar) -cv.createTrackbar(high_H_name, window_detection_name , high_H, max_value_H, on_high_H_thresh_trackbar) -cv.createTrackbar(low_S_name, window_detection_name , low_S, max_value, on_low_S_thresh_trackbar) -cv.createTrackbar(high_S_name, window_detection_name , high_S, max_value, on_high_S_thresh_trackbar) -cv.createTrackbar(low_V_name, window_detection_name , low_V, max_value, on_low_V_thresh_trackbar) -cv.createTrackbar(high_V_name, window_detection_name , high_V, max_value, on_high_V_thresh_trackbar) +cv.createTrackbar(low_H_name, window_detection_name, low_H, max_value_H, on_low_H_thresh_trackbar) +cv.createTrackbar(high_H_name, window_detection_name, high_H, max_value_H, on_high_H_thresh_trackbar) +cv.createTrackbar(low_S_name, window_detection_name, low_S, max_value, on_low_S_thresh_trackbar) +cv.createTrackbar(high_S_name, window_detection_name, high_S, max_value, on_high_S_thresh_trackbar) +cv.createTrackbar(low_V_name, window_detection_name, low_V, max_value, on_low_V_thresh_trackbar) +cv.createTrackbar(high_V_name, window_detection_name, high_V, max_value, on_high_V_thresh_trackbar) -#cv.createTrackbar('low_canny', 'canny', low_canny, 500, lcanny) +# cv.createTrackbar('low_canny', 'canny', low_canny, 500, lcanny) paused = False likelyGate = [] @@ -150,15 +172,15 @@ def allLarger(thresh): frame = untampered if ret: if not paused: - frame = cv.resize(frame, (0,0), fx=0.5, fy=0.5) + frame = cv.resize(frame, (0, 0), fx=0.5, fy=0.5) blur = cv.GaussianBlur(frame, (5, 5), 0) frame_HSV = cv.cvtColor(blur, cv.COLOR_BGR2HSV) - #frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) - #canny = cv.Canny(frame_gray, 0) - frame_threshold = cv.inRange(frame_HSV, (low_H, low_S, low_V), (high_H, high_S, high_V)) #low_S ideal = 98 - - #frame_threshold = cv.bitwise_not(frame_threshold) - res = cv.bitwise_and(frame,frame, mask= frame_threshold) + # frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) + # canny = cv.Canny(frame_gray, 0) + frame_threshold = cv.inRange(frame_HSV, (low_H, low_S, low_V), (high_H, high_S, high_V)) # low_S ideal = 98 + + # frame_threshold = cv.bitwise_not(frame_threshold) + res = cv.bitwise_and(frame, frame, mask=frame_threshold) res2, contours, hierarchy = cv.findContours(frame_threshold, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) contours.sort(key=heuristic, reverse=True) @@ -169,17 +191,17 @@ def allLarger(thresh): likelyGate = [{'cont': contours[0], 'heur': heur0}, {'cont': contours[1], 'heur': heur1}] untampered = np.copy(frame) if contours: - #likelyGate.append(contours[0]) - #findLikelyGate(likelyGate, contours) + # likelyGate.append(contours[0]) + # findLikelyGate(likelyGate, contours) drawRects(frame, likelyGate) cv.imshow(window_capture_name, frame) cv.imshow(window_detection_name, frame_threshold) - #cv.imshow('canny', canny) - + # cv.imshow('canny', canny) + key = cv.waitKey(30) if key == ord('q') or key == 27: break if key == ord('p'): paused = not paused -#generalized problem, giving center of object contrasting with water +# generalized problem, giving center of object contrasting with water diff --git a/perception/tasks/sanity_test.py b/perception/tasks/sanity_test.py index e3d07a3..8b22147 100644 --- a/perception/tasks/sanity_test.py +++ b/perception/tasks/sanity_test.py @@ -1,6 +1,8 @@ import multiprocessing import pytest +# TODO: integrate with pytests testing suite + def sanity_test(algorithm, test_imgs): """ Runs a sanity test on the algorithm that checks for run time and general exceptions. diff --git a/perception/tasks/segmentation/GateTaskExample.py.orig b/perception/tasks/segmentation/GateTaskExample.py.orig deleted file mode 100644 index 5563092..0000000 --- a/perception/tasks/segmentation/GateTaskExample.py.orig +++ /dev/null @@ -1,86 +0,0 @@ -from TaskPerceiver import TaskPerceiver -from typing import Tuple -from sys import argv as args -from combinedFilter import init_combined_filter -import numpy as np -import cv2 as cv -#from segmentation.aggregateRescaling import init_aggregate_rescaling - -class GateTask(TaskPerceiver): - def analyze(self, frame: np.ndarray, debug: bool) -> Tuple[float, float]: - """Takes in the background removed image and returns the center between - the two gate posts. - Args: - frame: The background removed frame to analyze - debug: Whether or not tot display intermediate images for debugging - - Returns: - (x,y) coordinate with center of gate - """ -<<<<<<< HEAD - filtered_frame_copies = [filtered_frame for _ in range[10]] - np.stack(filtered_frame_copies, axis = -1) - mask = cv.inRange(filtered_frame, np.array[190], ) - - filtered_frame = combined_filter(frame, display_figs=False) - if debug: - return ((250, 250), filtered_frame) -======= - filtered_frame = combined_filter(frame, display_figs=False) - filtered_frame_copies = [filtered_frame for _ in range(3)] - stacked_filter_frames = np.concatenate(filtered_frame_copies, axis = 2) - mask = cv.inRange(stacked_filter_frames, - np.array([100, 100, 100]), np.array([255, 255, 255])) - _, contours, _ = cv.findContours(mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) - if contours: - cnt = max(contours, key=self.findStraightness)#lambda x: cv.minAreaRect(x)[1][0] * cv.minAreaRect(x)[1][1]) - # sorted_straight = sorted(contours, key=self.findStraightness) - # sorted_size = sorted(contours, key=cv.contourArea) - #todo: use these sorted lists and weights to each value to give two best values - rect = cv.minAreaRect(cnt) - boxpts = cv.boxPoints(rect) - box = np.int0(boxpts) - cv.drawContours(stacked_filter_frames,[box],0,(0,0,255),5) - for corner in boxpts: - cv.circle(stacked_filter_frames, (corner[0], corner[1]), 10, (0,0,255), -1) - - if debug: - return ((250, 250), stacked_filter_frames) ->>>>>>> origin/gate-task-example - return (250, 250) - - def findStraightness(self, contour): # output number = contour area/convex area, the bigger the straightest - hull = cv.convexHull(contour, False) - contour_area = cv.contourArea(contour) - hull_area = cv.contourArea(hull) - return 10 * contour_area + 3 * (hull_area - contour_area) - -# this part is temporary and will be covered by other files in the future -if __name__ == '__main__': - combined_filter = init_combined_filter() - cap = cv.VideoCapture(args[1]) - ret_tries = True - gate_task = GateTask() - once = False - while 1 and ret_tries < 50: - ret, frame = cap.read() - if ret: - frame = cv.resize(frame, None, fx=0.4, fy=0.4) - - - ### FUNCTION CALL, can change this - (x, y), filtered_frame = gate_task.analyze(frame, True) - cv.putText(frame, "x: %.2f" % x + " y: %.2f" % y, - (20, frame.shape[0] - 20), cv.FONT_HERSHEY_SIMPLEX, - 2.0, (0, 165, 255), 3) - cv.imshow('original', frame) - cv.imshow('filtered_frame', filtered_frame) - if not once: - print(filtered_frame) - once = True - ret_tries = 0 - k = cv.waitKey(60) & 0xff - if k == 27: - break - else: - ret_tries += 1 diff --git a/perception/tasks/segmentation/__init__.py b/perception/tasks/segmentation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/perception/tasks/segmentation/aggregateRescaling.py b/perception/tasks/segmentation/aggregateRescaling.py index 3e2d9ea..e0f4354 100644 --- a/perception/tasks/segmentation/aggregateRescaling.py +++ b/perception/tasks/segmentation/aggregateRescaling.py @@ -3,6 +3,8 @@ import numpy as np import numpy.linalg as LA + +# TODO: port to vis + TaskPerciever format or remove # Jenny -> unsigned ints fixed the problem # Damas -> flip weight vector every frame @@ -42,9 +44,9 @@ def aggregate_rescaling(frame): # you only pca once red -= max_min['min'] red *= 255.0 / (max_min['max'] - max_min['min']) """ - if False:#not paused: - print(np.min(red), np.max(red), max_min['min'], max_min['max']) - """ + if False:#not paused: + print(np.min(red), np.max(red), max_min['min'], max_min['max']) + """ red = red.astype(np.uint8) red = np.expand_dims(red, axis=2) red = np.concatenate((red, red, red), axis=2) diff --git a/perception/tasks/segmentation/combinedFilter.py b/perception/tasks/segmentation/combinedFilter.py index 22977d1..ec32624 100644 --- a/perception/tasks/segmentation/combinedFilter.py +++ b/perception/tasks/segmentation/combinedFilter.py @@ -1,6 +1,8 @@ import cv2 import numpy as np +# TODO: port to vis + TaskPerciever format or remove + from sys import argv as args from perception.tasks.segmentation.aggregateRescaling import init_aggregate_rescaling from perception.tasks.segmentation.peak_removal_adaptive_thresholding import filter_out_highest_peak_multidim @@ -23,7 +25,7 @@ def combined_filter(frame, custom_weights=None, display_figs=False, print_weight custom_weights=custom_weights, print_weights=print_weights) - other_frame = other_frame[:,:,:1] + other_frame = other_frame[:, :, :1] if display_figs: cv2.imshow('original', frame) diff --git a/perception/tasks/segmentation/kmeans.py b/perception/tasks/segmentation/kmeans.py new file mode 100644 index 0000000..507b261 --- /dev/null +++ b/perception/tasks/segmentation/kmeans.py @@ -0,0 +1,99 @@ +import cv2 +import numpy as np +import matplotlib.pyplot as plt +from scipy.signal import find_peaks, peak_widths +from sys import argv as args + +# TODO: port to vis + TaskPerciever format or remove + +######################################################################## +# An attempt at an adaptive thresholding algorithm based on the frequency +# of pixel values ("peaks" if looking at a histogram of # pixels vs pixel value of a frame) +# +# *1. *** best of the three *** filter_out_highest_peak_multidim +# pools together how "peak-like" each pixel is in all of the color channels +# of the frame to make a final decision on what is the background +# 2. init_filter_out_highest_peak +# gets rid of large peaks in many different color channels individually +# 3. remove_blotchy_chunks +# places a mask over areas that have lots of edges, which in many cases +# is equivalent to places with lots of noise +######################################################################## + + +def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): + """ Attempts to use kmeans to segment the frame into num_group features + (not including the background), denoted by a very large value in votes. + votes is an output of the filter_out_highest_peak_multidim() function + Output: frame_shape x num_groups 3D matrix. Get a group mapped to the + frame by doing groups[:,:,group_num] """ + votes = np.float32(votes).flatten() + + # Make kmeans only consider the non-background pixels + background = np.zeros(votes.shape) + background[votes >= np.percentile(votes, percentile)] = 1 + cluster_data = votes[background == 0] + cluster_indexes = np.array(range(len(votes)))[background == 0] + + # Do kmeans + criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) + flags = cv2.KMEANS_RANDOM_CENTERS + compactness, labels, centers = cv2.kmeans(cluster_data, num_groups, None, criteria, 10, flags) + + # Reconstruct the original votes array with background's label = -1 + label_arr = np.empty(votes.shape) + label_arr[background == 1] = -1 + for i in range(num_groups): + label_arr[cluster_indexes[labels.flatten() == i]] = i + + unique_labels, label_counts = np.unique(label_arr, return_counts=True) + label_order = list(range(np.int0(np.amax(unique_labels)) + 2)) # something is erroring here + if len(label_counts) < num_groups + 1: + # add in a slot for the background if no background is found + label_counts = np.insert(label_counts, 0, 0) + + label_order.sort(key=lambda x: label_counts[x]) + + groups = np.empty((frame_shape[0], frame_shape[1], num_groups + 1)) + for i, l in enumerate(label_order): + group = np.zeros(votes.shape) + group[label_arr.flatten() == l - 1] = 255 + groups[:, :, i] = np.reshape(group, frame_shape[:2]) + + # for i in range(len(unique_labels)): + # cv2.imshow(str(i) + " label", groups[:,:,i]) + + return groups + + +########################################### +# Main Body +########################################### + +if __name__ == "__main__": + # For testing porpoises + cap = cv2.VideoCapture(args[1]) + ret, frame = cap.read() + out = cv2.VideoWriter('out.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30.0, + (int(frame.shape[1] * 0.4), int(frame.shape[0] * 0.4))) + + ret_tries = 0 + + while (1 and ret_tries < 50): + ret, frame = cap.read() + + if ret: + frame = cv2.resize(frame, None, fx=0.4, fy=0.4) + + cv2.imshow('original', frame) + plt.pause(0.001) + + ret_tries = 0 + k = cv2.waitKey(60) & 0xff + if k == 27: + break + else: + ret_tries += 1 + cv2.destroyAllWindows() + cap.release() + out.release() \ No newline at end of file diff --git a/perception/tasks/segmentation/peak_removal_adaptive_thresholding.py b/perception/tasks/segmentation/peak_removal_adaptive_thresholding.py index 892dc02..d76a17b 100644 --- a/perception/tasks/segmentation/peak_removal_adaptive_thresholding.py +++ b/perception/tasks/segmentation/peak_removal_adaptive_thresholding.py @@ -3,6 +3,9 @@ import matplotlib.pyplot as plt from scipy.signal import find_peaks, peak_widths from sys import argv as args + +# TODO: port to vis + TaskPerciever format or remove + ######################################################################## # An attempt at an adaptive thresholding algorithm based on the frequency # of pixel values ("peaks" if looking at a histogram of # pixels vs pixel value of a frame) @@ -51,6 +54,7 @@ # thresholds_used = [h_low, s_low, v_low, h_hi, s_hi, v_hi] + def init_test_hsv_thresholds(thresholds): # Keep track of previous threhold values to see if the user is using the trackbar # is there a function that detects whether the mouse button is down? @@ -61,25 +65,25 @@ def nothing(x): pass cv2.namedWindow('ideal thresholding') - cv2.createTrackbar('h_low','ideal thresholding',h_low,255,nothing) - cv2.createTrackbar('s_low','ideal thresholding',s_low,255,nothing) - cv2.createTrackbar('v_low','ideal thresholding',v_low,255,nothing) - cv2.createTrackbar('h_high','ideal thresholding',h_hi,255,nothing) - cv2.createTrackbar('s_high','ideal thresholding',s_hi,255,nothing) - cv2.createTrackbar('v_high','ideal thresholding',v_hi,255,nothing) + cv2.createTrackbar('h_low', 'ideal thresholding', h_low, 255, nothing) + cv2.createTrackbar('s_low', 'ideal thresholding', s_low, 255, nothing) + cv2.createTrackbar('v_low', 'ideal thresholding', v_low, 255, nothing) + cv2.createTrackbar('h_high', 'ideal thresholding', h_hi, 255, nothing) + cv2.createTrackbar('s_high', 'ideal thresholding', s_hi, 255, nothing) + cv2.createTrackbar('v_high', 'ideal thresholding', v_hi, 255, nothing) def test_hsv_thresholds(frame, thresholds): nonlocal prev_h_low, prev_s_low, prev_v_low, prev_h_hi, prev_s_hi, prev_v_hi - h_low_track = cv2.getTrackbarPos('h_low','ideal thresholding') - s_low_track = cv2.getTrackbarPos('s_low','ideal thresholding') - v_low_track = cv2.getTrackbarPos('v_low','ideal thresholding') - h_hi_track = cv2.getTrackbarPos('h_high','ideal thresholding') - s_hi_track = cv2.getTrackbarPos('s_high','ideal thresholding') - v_hi_track = cv2.getTrackbarPos('v_high','ideal thresholding') + h_low_track = cv2.getTrackbarPos('h_low', 'ideal thresholding') + s_low_track = cv2.getTrackbarPos('s_low', 'ideal thresholding') + v_low_track = cv2.getTrackbarPos('v_low', 'ideal thresholding') + h_hi_track = cv2.getTrackbarPos('h_high', 'ideal thresholding') + s_hi_track = cv2.getTrackbarPos('s_high', 'ideal thresholding') + v_hi_track = cv2.getTrackbarPos('v_high', 'ideal thresholding') - if h_low_track!=prev_h_low or s_low_track!=prev_s_low or v_low_track!=prev_v_low \ - or h_hi_track!=prev_h_hi or s_hi_track!=prev_s_hi or v_hi_track!=prev_v_hi: + if h_low_track != prev_h_low or s_low_track != prev_s_low or v_low_track != prev_v_low \ + or h_hi_track != prev_h_hi or s_hi_track != prev_s_hi or v_hi_track != prev_v_hi: # If user is adjusting the trackbars, use the user input thresholds_used = [h_low_track, s_low_track, v_low_track, h_hi_track, s_hi_track, v_hi_track] else: @@ -92,9 +96,9 @@ def test_hsv_thresholds(frame, thresholds): cv2.setTrackbarPos('s_high', 'ideal thresholding', thresholds_used[4]) cv2.setTrackbarPos('v_high', 'ideal thresholding', thresholds_used[5]) - hsv = cv2.cvtColor(frame,cv2.COLOR_BGR2HSV) + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) mask = cv2.inRange(hsv, np.array(thresholds_used[:3]), np.array(thresholds_used[3:])) - res = cv2.bitwise_and(frame,frame, mask= mask) + res = cv2.bitwise_and(frame, frame, mask=mask) cv2.imshow('ideal thresholding', res) @@ -103,20 +107,22 @@ def test_hsv_thresholds(frame, thresholds): return test_hsv_thresholds + def hsv_threshold(frame, thresh_used): - hsv = cv2.cvtColor(frame,cv2.COLOR_BGR2HSV) + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) mask = cv2.inRange(hsv, np.array(thresh_used[:3]), np.array(thresh_used[3:])) - res = cv2.bitwise_and(frame,frame, mask= mask) + res = cv2.bitwise_and(frame, frame, mask=mask) return thresh_used, res + def disp_hist(frame, title, labels, colors): - frame0 = frame[:,:,0].flatten() + frame0 = frame[:, :, 0].flatten() frame0 = frame0[frame0 > 0] - frame1 = frame[:,:,1].flatten() + frame1 = frame[:, :, 1].flatten() frame1 = frame1[frame1 > 0] - frame2 = frame[:,:,2].flatten() + frame2 = frame[:, :, 2].flatten() frame2 = frame2[frame2 > 0] plt.figure(hash(title)) @@ -131,6 +137,7 @@ def disp_hist(frame, title, labels, colors): plt.legend() plt.draw() + def find_peak_ranges(frame, display_plots=False, title=None, labels=None, colors=None): """ Finds a returns the widest peak's x-range in all three channels of frame Result is formatted to fit cv2.inRange() -> ((low1, low2, low3), (hi1, hi2, hi3)) @@ -138,8 +145,9 @@ def find_peak_ranges(frame, display_plots=False, title=None, labels=None, colors # TODO: Maybe use a different combination of peak characteristics to more accurately # select the entire peak (only the tip is selected right now) - peak_width_height = 0.95 # How far down the peak that the algorithm draws - # the horizontal width line + peak_width_height = 0.95 # How far down the peak that the algorithm draws + + # the horizontal width line def find_highest_peak(channel, display_plots=False): """ Finds and returns the x-range of the highest peak in the @@ -148,33 +156,34 @@ def find_highest_peak(channel, display_plots=False): f = frame[:, :, channel].flatten() # Some semi-hardcoded values :) - num_bins = max(int((np.amax(f)-np.amin(f)) / 4), 10) + num_bins = max(int((np.amax(f) - np.amin(f)) / 4), 10) # num_bins = 30 hist, bins = np.histogram(f, bins=num_bins) - hist[0] = 0 # get rid of stuff that was thresholded to 0 - hist = np.hstack([hist, [0]]) # make stuff at 255 into a peak - bins = np.hstack([bins, [bins[bins.shape[0]-1] + 1]]) + hist[0] = 0 # get rid of stuff that was thresholded to 0 + hist = np.hstack([hist, [0]]) # make stuff at 255 into a peak + bins = np.hstack([bins, [bins[bins.shape[0] - 1] + 1]]) peaks, properties = find_peaks(hist, height=0.1) if len(peaks) > 0: i = np.argmax(properties['peak_heights']) widths = peak_widths(hist, peaks, rel_height=peak_width_height)[0] # i = np.argmax(widths) - largest_peak = (int((bins[peaks[i]]+bins[peaks[i]+1])//2-widths[i]//2), - int((bins[peaks[i]]+bins[peaks[i]+1])//2+widths[i]//2)) # beginning and end of the peak + largest_peak = (int((bins[peaks[i]] + bins[peaks[i] + 1]) // 2 - widths[i] // 2), + int((bins[peaks[i]] + bins[peaks[i] + 1]) // 2 + widths[ + i] // 2)) # beginning and end of the peak if display_plots: ax = plt.gca() print(max(f)) ax.set_xlim([0, max(255, max(f))]) - #Plot values in this channel - plt.plot(bins[1:],hist, label=labels[channel], color=colors[channel]) + # Plot values in this channel + plt.plot(bins[1:], hist, label=labels[channel], color=colors[channel]) # Plot peaks - plt.plot(bins[peaks+1], hist[peaks], "x") + plt.plot(bins[peaks + 1], hist[peaks], "x") # Plot peak widths - plt.hlines(hist[peaks]*0.9, bins[peaks+1]-widths//2, bins[peaks+1]+widths//2) + plt.hlines(hist[peaks] * 0.9, bins[peaks + 1] - widths // 2, bins[peaks + 1] + widths // 2) else: largest_peak = (0, 0) @@ -184,7 +193,7 @@ def find_highest_peak(channel, display_plots=False): fig = plt.figure(hash(title)) plt.clf() - background = (np.empty(frame.shape[2]),np.empty(frame.shape[2])) + background = (np.empty(frame.shape[2]), np.empty(frame.shape[2])) for channel in range(frame.shape[2]): low, high = find_highest_peak(channel, display_plots) background[0][channel] = low @@ -197,12 +206,14 @@ def find_highest_peak(channel, display_plots=False): return background + def plot_peaks(frame, title, labels, colors): # Shh this is just a helper function that makes the code more readable # Not to be used in practice. # NOTE: you need to call plt.pause(0.001) afterwards to render the plot find_peak_ranges(frame, True, title, labels, colors) + def init_filter_out_highest_peak(filters, return_colorspace="any", input_colorspace="bgr"): """ Takes in an hsv image! Returns an hsv image""" # low pass filter @@ -210,8 +221,8 @@ def init_filter_out_highest_peak(filters, return_colorspace="any", input_colorsp # lambda = 0.9-0.4 prev_hsv_threshes = [[] for i in range(len(filters))] - hsv_labels = (('H','S','V'), ("red","purple","gray")) - bgr_labels = (('B','G','R'), ("blue","green","red")) + hsv_labels = (('H', 'S', 'V'), ("red", "purple", "gray")) + bgr_labels = (('B', 'G', 'R'), ("blue", "green", "red")) # Figure out how the procedure to convert among hsv and bgr. # Format of stuff in fitler_fns: @@ -220,9 +231,9 @@ def init_filter_out_highest_peak(filters, return_colorspace="any", input_colorsp curr_color = input_colorspace for f in filters: if f != curr_color: - filter_fns.append(['c',f]) + filter_fns.append(['c', f]) curr_color = f - filter_fns.append(['f',f]) + filter_fns.append(['f', f]) if return_colorspace != "any" and return_colorspace != curr_color: filter_fns.append(['c', return_colorspace]) @@ -236,27 +247,27 @@ def filter_out_highest_peak(frame, cache, display_plots=False, title=None, label # calculate average for i in range(2): for j in range(3): - background_thresh[i][j] = (background_thresh[i][j] + sum([c[i][j] for c in cache])) // (len(cache) + 1) + background_thresh[i][j] = (background_thresh[i][j] + sum([c[i][j] for c in cache])) // ( + len(cache) + 1) background_mask = cv2.bitwise_not(cv2.bitwise_or( - cv2.inRange(frame[:, :, 0], background_thresh[0][0], background_thresh[1][0]), - cv2.inRange(frame[:, :, 1], background_thresh[0][1], background_thresh[1][1]), - cv2.inRange(frame[:, :, 2], background_thresh[0][2], background_thresh[1][2]) - )) - no_background = cv2.bitwise_and(frame,frame, mask=background_mask) + cv2.inRange(frame[:, :, 0], background_thresh[0][0], background_thresh[1][0]), + cv2.inRange(frame[:, :, 1], background_thresh[0][1], background_thresh[1][1]), + cv2.inRange(frame[:, :, 2], background_thresh[0][2], background_thresh[1][2]) + )) + no_background = cv2.bitwise_and(frame, frame, mask=background_mask) return background_thresh, raw_thresh, no_background def combine_threshes(th1, th2): - return ([min(th1[0][0], th2[0][0]), min(th1[0][1], th2[0][1]), min(th1[0][2], th2[0][2])], - [max(th1[1][0], th2[1][0]), max(th1[1][1], th2[1][1]), max(th1[1][2], th2[1][2])]) + return ([min(th1[0][0], th2[0][0]), min(th1[0][1], th2[0][1]), min(th1[0][2], th2[0][2])], + [max(th1[1][0], th2[1][0]), max(th1[1][1], th2[1][1]), max(th1[1][2], th2[1][2])]) def bgr_thresh2hsv_thresh(th): th = cv2.cvtColor(np.array([[th[0]], [th[1]]], np.uint8), cv2.COLOR_BGR2HSV).tolist() return ([min(th[0][0][0], th[1][0][0]), min(th[0][0][1], th[1][0][1]), min(th[0][0][2], th[1][0][2])], [max(th[0][0][0], th[1][0][0]), max(th[0][0][1], th[1][0][1]), max(th[0][0][2], th[1][0][2])]) - def do_filter(frame, display_plots=False): nonlocal prev_hsv_threshes if len(prev_hsv_threshes[0]) == lpf_cache_size: @@ -298,32 +309,33 @@ def do_filter(frame, display_plots=False): return do_filter + def keep_highest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, title=None, label='1', color='blue'): """ Returns a mask for the frame that keeps the num_peaks highest peaks in the histogram of pixel values. Only works for grayscale/1-channel images (to speed this up) Shape of frame must have 3 dimensions (pass in np.expand_dims(frame, 2) if erroring) """ # Some semi-thresholded values :) - num_bins = max(int((np.amax(frame)-np.amin(frame)) / 4), 10) + num_bins = max(int((np.amax(frame) - np.amin(frame)) / 4), 10) hist, bins = np.histogram(frame, bins=num_bins) - hist[0] = 0 # get rid of stuff that was thresholded to 0 - hist = np.hstack([hist, [0]]) # make stuff at 255 into a peak - bins = np.hstack([bins, [bins[bins.shape[0]-1] + 1]]) + hist[0] = 0 # get rid of stuff that was thresholded to 0 + hist = np.hstack([hist, [0]]) # make stuff at 255 into a peak + bins = np.hstack([bins, [bins[bins.shape[0] - 1] + 1]]) peaks, properties = find_peaks(hist, prominence=100) widths = peak_widths(hist, peaks, rel_height=peak_width_height)[0] if len(peaks) > 0: i = len(peaks) - 1 - mask = cv2.inRange(frame, (bins[peaks[i]]+bins[peaks[i]+1])//2-widths[i]*2, - (bins[peaks[i]]+bins[peaks[i]+1])//2+widths[i]*2) + mask = cv2.inRange(frame, (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 - widths[i] * 2, + (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 + widths[i] * 2) # To support keeping multiple peaks for j in range(num_peaks - 1): i = len(peaks) - 2 - j if i >= 0: - mask = cv2.bitwise_or(cv2.inRange(frame, (bins[peaks[i]]+bins[peaks[i]+1])//2-widths[i], - (bins[peaks[i]]+bins[peaks[i]+1])//2+widths[i]), mask) + mask = cv2.bitwise_or(cv2.inRange(frame, (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 - widths[i], + (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 + widths[i]), mask) # frame = cv2.bitwise_and(frame, frame, mask=mask) else: mask = np.ones(frame.shape, np.uint8) @@ -334,12 +346,12 @@ def keep_highest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, titl ax = plt.gca() ax.set_xlim([0, 255]) - #Plot values in this channel - plt.plot(bins[1:],hist, label=label, color=color) + # Plot values in this channel + plt.plot(bins[1:], hist, label=label, color=color) # Plot peaks - plt.plot(bins[peaks+1], hist[peaks], "x") + plt.plot(bins[peaks + 1], hist[peaks], "x") # Plot peak widths - plt.hlines(hist[peaks]*0.9, bins[peaks+1]-widths//2, bins[peaks+1]+widths//2) + plt.hlines(hist[peaks] * 0.9, bins[peaks + 1] - widths // 2, bins[peaks + 1] + widths // 2) plt.title(title) plt.legend() @@ -347,31 +359,32 @@ def keep_highest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, titl return mask + def delete_lowest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, title=None, label='1', color='blue'): """ Returns a mask for the frame that deletes the num_peaks lowest-valued peaks in the histogram of pixel values. Only works for grayscale/1-channel images (to speed this up) """ # Some semi-thresholded values :) - num_bins = max(int((np.amax(frame)-np.amin(frame)) / 4), 10) + num_bins = max(int((np.amax(frame) - np.amin(frame)) / 4), 10) hist, bins = np.histogram(frame, bins=num_bins) - hist[0] = 0 # get rid of stuff that was thresholded to 0 + hist[0] = 0 # get rid of stuff that was thresholded to 0 peaks, properties = find_peaks(hist, prominence=100) widths = peak_widths(hist, peaks, rel_height=peak_width_height)[0] if len(peaks) > 0: i = 0 - mask = cv2.bitwise_not(cv2.inRange(frame, (bins[peaks[i]]+bins[peaks[i]+1])//2-widths[i]*2, - (bins[peaks[i]]+bins[peaks[i]+1])//2+widths[i]*2)) + mask = cv2.bitwise_not(cv2.inRange(frame, (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 - widths[i] * 2, + (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 + widths[i] * 2)) # To support deleting multiple peaks for j in range(num_peaks - 1): i = j + 1 if len(peaks) > i: mask = cv2.bitwise_and(cv2.bitwise_not(cv2.inRange( - frame, (bins[peaks[i]]+bins[peaks[i]+1])//2-widths[i]*2, - (bins[peaks[i]]+bins[peaks[i]+1])//2+widths[i]*2)), mask) + frame, (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 - widths[i] * 2, + (bins[peaks[i]] + bins[peaks[i] + 1]) // 2 + widths[i] * 2)), mask) else: mask = np.ones(frame.shape, np.uint8) # frame = cv2.bitwise_and(frame, frame, mask=mask) @@ -382,12 +395,12 @@ def delete_lowest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, tit ax = plt.gca() ax.set_xlim([0, 255]) - #Plot values in this channel - plt.plot(bins[1:],hist, label=label, color=color) + # Plot values in this channel + plt.plot(bins[1:], hist, label=label, color=color) # Plot peaks - plt.plot(bins[peaks+1], hist[peaks], "x") + plt.plot(bins[peaks + 1], hist[peaks], "x") # Plot peak widths - plt.hlines(hist[peaks]*0.9, bins[peaks+1]-widths//2, bins[peaks+1]+widths//2) + plt.hlines(hist[peaks] * 0.9, bins[peaks + 1] - widths // 2, bins[peaks + 1] + widths // 2) plt.title(title) plt.legend() @@ -395,6 +408,7 @@ def delete_lowest_valued_peaks_mask(frame, num_peaks=1, display_plots=False, tit return mask + def remove_blotchy_chunks(frame, kernel_size=201, iterations=1, display_imgs=False): """ Works best when object isn't surrounded by blotchy stuff """ edges = cv2.Canny(frame, 100, 150) @@ -416,6 +430,7 @@ def remove_blotchy_chunks(frame, kernel_size=201, iterations=1, display_imgs=Fal return result + def filter_out_highest_peak_multidim(frame, res=69, percentile=10, custom_weights=None, print_weights=False): """ Estimates the "peak-ness" of each pixel in frame across color channels and thresholds out pixels that were "peak-like" in many colorspaces. @@ -444,7 +459,7 @@ def get_peak_votes(frame): if res == 1: vote_arr = dist[frame] else: - dist = np.array([np.mean(dist[i*res:i*res+res]) for i in range(len(dist) // res + 1)]) + dist = np.array([np.mean(dist[i * res:i * res + res]) for i in range(len(dist) // res + 1)]) vote_arr = dist[frame // res] return recommended_weight, vote_arr @@ -455,7 +470,7 @@ def get_peak_votes(frame): if print_weights: print('------------------------', custom_weights) for ch in range(frame.shape[2]): - weight, vote_arr = get_peak_votes(frame[:,:,ch]) + weight, vote_arr = get_peak_votes(frame[:, :, ch]) if custom_weights is not None: weight = custom_weights[ch] if print_weights: @@ -472,6 +487,7 @@ def get_peak_votes(frame): return overall_votes, cv2.bitwise_and(frame, frame, mask=overall_mask) + def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): """ Attempts to use kmeans to segment the frame into num_group features (not including the background), denoted by a very large value in votes. @@ -482,23 +498,23 @@ def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): # Make kmeans only consider the non-background pixels background = np.zeros(votes.shape) - background[votes>=np.percentile(votes, percentile)] = 1 - cluster_data = votes[background==0] - cluster_indexes = np.array(range(len(votes)))[background==0] + background[votes >= np.percentile(votes, percentile)] = 1 + cluster_data = votes[background == 0] + cluster_indexes = np.array(range(len(votes)))[background == 0] # Do kmeans criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) flags = cv2.KMEANS_RANDOM_CENTERS - compactness,labels,centers = cv2.kmeans(cluster_data,num_groups,None,criteria,10,flags) + compactness, labels, centers = cv2.kmeans(cluster_data, num_groups, None, criteria, 10, flags) # Reconstruct the original votes array with background's label = -1 label_arr = np.empty(votes.shape) - label_arr[background==1] = -1 + label_arr[background == 1] = -1 for i in range(num_groups): - label_arr[cluster_indexes[labels.flatten()==i]] = i + label_arr[cluster_indexes[labels.flatten() == i]] = i unique_labels, label_counts = np.unique(label_arr, return_counts=True) - label_order = list(range(np.int0(np.amax(unique_labels)) + 2)) # something is erroring here + label_order = list(range(np.int0(np.amax(unique_labels)) + 2)) # something is erroring here if len(label_counts) < num_groups + 1: # add in a slot for the background if no background is found label_counts = np.insert(label_counts, 0, 0) @@ -508,14 +524,15 @@ def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): groups = np.empty((frame_shape[0], frame_shape[1], num_groups + 1)) for i, l in enumerate(label_order): group = np.zeros(votes.shape) - group[label_arr.flatten()==l - 1] = 255 - groups[:,:,i] = np.reshape(group, frame_shape[:2]) + group[label_arr.flatten() == l - 1] = 255 + groups[:, :, i] = np.reshape(group, frame_shape[:2]) # for i in range(len(unique_labels)): # cv2.imshow(str(i) + " label", groups[:,:,i]) return groups + ########################################### # Main Body ########################################### @@ -524,7 +541,8 @@ def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): # For testing porpoises cap = cv2.VideoCapture(args[1]) ret, frame = cap.read() - out = cv2.VideoWriter('out.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30.0, (int(frame.shape[1]*0.4), int(frame.shape[0]*0.4))) + out = cv2.VideoWriter('out.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30.0, + (int(frame.shape[1] * 0.4), int(frame.shape[0] * 0.4))) if testing: test_hsv_thresholds = init_test_hsv_thresholds(thresholds_used) @@ -539,8 +557,10 @@ def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): if ret: frame = cv2.resize(frame, None, fx=0.4, fy=0.4) - votes, multi_filter1 = filter_out_highest_peak_multidim(np.dstack([frame, cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)]), custom_weights=[1, 1, 1, 1, 1, 1]) - votes, multi_filter2 = filter_out_highest_peak_multidim(np.dstack([frame, cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)])) + votes, multi_filter1 = filter_out_highest_peak_multidim( + np.dstack([frame, cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)]), custom_weights=[1, 1, 1, 1, 1, 1]) + votes, multi_filter2 = filter_out_highest_peak_multidim( + np.dstack([frame, cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)])) multi_filter1 = multi_filter1[:, :, :3] multi_filter2 = multi_filter2[:, :, :3] @@ -567,4 +587,4 @@ def k_means_segmentation(votes, frame_shape, num_groups=2, percentile=10): ret_tries += 1 cv2.destroyAllWindows() cap.release() - out.release() \ No newline at end of file + out.release() diff --git a/perception/tasks/slots/__init__.py b/perception/tasks/slots/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/perception/tasks/path_marker/play_slots_detection.py b/perception/tasks/slots/play_slots_detection.py similarity index 99% rename from perception/tasks/path_marker/play_slots_detection.py rename to perception/tasks/slots/play_slots_detection.py index 0ebb4b5..4b8b741 100644 --- a/perception/tasks/path_marker/play_slots_detection.py +++ b/perception/tasks/slots/play_slots_detection.py @@ -2,6 +2,8 @@ import cv2 import sys +# TODO: port to vis + TaskPerciever format or remove + #### TODO: maybe look into pattern matching # Data fron the new course footage dropbox folder diff --git a/perception/tasks/spinny/spinny_wheel_detection.py b/perception/tasks/spinny/spinny_wheel_detection.py deleted file mode 100644 index 8f0f72f..0000000 --- a/perception/tasks/spinny/spinny_wheel_detection.py +++ /dev/null @@ -1,112 +0,0 @@ -import numpy as np -import cv2 -import argparse -import sys -import time - -#CHANGE PARAMETER IN CALL TO "THRESH" FUNCTION TO CHANGE COLOR -file_name = "GOPR1145.MP4" #video file from dropbox -vid = cv2.VideoCapture(file_name) -frames = 0 -avgLength = 10 -centers = [] - -def thresh(frame, color='red'): - - blur = cv2.GaussianBlur(frame, (5, 5), 0) - hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV) - - if color == 'red': - lower = np.uint8([29,77,36]) - upper = np.uint8([130,250,255]) - mask = cv2.inRange(hsv,lower,upper) - mask = cv2.bitwise_not(mask) - elif color == 'blue': - lower = np.uint8([86,141,0]) - upper = np.uint8([106,220,168]) - mask = cv2.inRange(hsv,lower,upper) - else: - lower = np.uint8([66,208,157]) - upper = np.uint8([86,255,209]) - mask = cv2.inRange(hsv,lower,upper) - - return mask - -def heuristic(contour): - rect = cv2.minAreaRect(contour) - area = rect[1][0] * rect[1][1] - diff = cv2.contourArea(cv2.convexHull(contour)) - cv2.contourArea(contour) - cent = rect[0] - dist = 0 - if len(likelySection) > 1 and allLarger(60): - cen0 = cv2.minAreaRect(likelySection[0]['cont'])[0] - dis0 = np.linalg.norm(np.array(cent) - np.array(cen0)) - cen1 = cv2.minAreaRect(likelySection[1]['cont'])[0] - dis1 = np.linalg.norm(np.array(cent) - np.array(cen1)) - dist = min([dis0, dis1]) - heur = area - 3 * diff - 20 * dist - return heur - -def allLarger(thresh): - for cnt in likelySection: - if cnt['heur'] < thresh: - return False - return True - -def drawRects(frame, contours): - tempPts = [] - for cnt in contours: - rect = cv2.minAreaRect(cnt['cont']) - boxpts = cv2.boxPoints(rect) - box = np.int0(boxpts) - cv2.drawContours(frame,[box],0,(0,0,255),1) - cv2.drawContours(frame, [cnt['cont']],0,(0,255,0),1) - cv2.drawContours(frame, [cv2.convexHull(cnt['cont'])],0,(255,0,0),1) - tempPts.append(rect[0]) - #cv2.putText(frame, str(cnt['heur']), (int(rect[0][0]), int(rect[0][1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255)) - if len(tempPts) > 1 and allLarger(60): - cv2.circle(frame, (int(tempPts[0][0]), int(tempPts[0][1])), 10, (0,0,255), -1) - cv2.circle(frame, (int(tempPts[1][0]), int(tempPts[1][1])), 10, (0,0,255), -1) -def midPt(pt1, pt2): - return ((pt1[0] + pt2[0]) / 2, (pt1[1] + pt2[1]) / 2) - -def getAvgPt(pt): - points.append(pt) - exes = list(map(lambda x: x[0], points)) - whys = list(map(lambda y: y[1], points)) - - if len(points) > 50: - del points[:10] - return (int(sum(exes) / len(exes)), int(sum(whys) / len(whys))) - -likelySection = [] -points = [] -while vid.isOpened(): - start = time.time() - ret, frame = vid.read() - if(ret == False): - continue - - threshed = thresh(frame,'other') - res, contours, hierarchy = cv2.findContours(threshed, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - - contours.sort(key=heuristic, reverse = True) - if len(contours) > 1: - c1 = contours[0] - c2 = contours[1] - heur0 = heuristic(c1) - heur1 = heuristic(c2) - likelySection = [{'cont': contours[0], 'heur': heur0}, {'cont': contours[1], 'heur': heur1}] - untampered = np.copy(frame) - if contours: - drawRects(frame, likelySection) - cv2.imshow("Frame", frame) - cv2.imshow('Res', res) - - if (cv2.waitKey(1) & 0xFF) == ord('q') or frames > 900: - break - -vid.release() -cv2.destroyAllWindows() - - diff --git a/perception/vis/FrameWrapper.py b/perception/vis/FrameWrapper.py index 4eb6e01..fe751d1 100644 --- a/perception/vis/FrameWrapper.py +++ b/perception/vis/FrameWrapper.py @@ -24,7 +24,7 @@ class FrameWrapper(): WEBCAM_TRIES = 10 def __init__(self, filenames, resize=1): - self.filenames = filenames # Get this list of relative paths to files from vis + self.filenames = filenames # Get this list of relative paths to files from vis # There aren't any checks for resize==1 to improve speed b/c this expects resize != 1 self.resize = resize @@ -46,27 +46,24 @@ def __next__(self): ret, frame = self.next_data[1].read() if ret: return cv2.resize(frame, None, fx=self.resize, fy=self.resize) - else: - print("WARNING: Failed to get frame from video {}. Try {}." \ - .format(self.filenames[self.index], i), file=sys.stderr) + print("WARNING: Failed to get frame from video {}. Try {}." \ + .format(self.filenames[self.index], i), file=sys.stderr) self.next_data_obj() elif self.next_data[0] == "i": # Image img = self.next_data[1] self.next_data_obj() if img is not None: return cv2.resize(img, None, fx=self.resize, fy=self.resize) - else: - print("WARNING: Failed to get image {}." \ - .format(self.filenames[self.index-1]), file=sys.stderr) + print("WARNING: Failed to get image {}." \ + .format(self.filenames[self.index-1]), file=sys.stderr) else: # Webcam # Try to get a frame out at most WEBCAM_TRIES times. for i in range(self.WEBCAM_TRIES): ret, frame = self.next_data[1].read() if ret: return cv2.resize(frame, None, fx=self.resize, fy=self.resize) - else: - print("WARNING: Failed to get frame from webcam. Try {}." \ - .format(i), file=sys.stderr) + print("WARNING: Failed to get frame from webcam. Try {}." \ + .format(i), file=sys.stderr) self.next_data_obj() raise StopIteration diff --git a/perception/tasks/TestTasks/TestAlgo.py b/perception/vis/TestAlgo.py similarity index 54% rename from perception/tasks/TestTasks/TestAlgo.py rename to perception/vis/TestAlgo.py index 3410aad..3d2ebac 100644 --- a/perception/tasks/TestTasks/TestAlgo.py +++ b/perception/vis/TestAlgo.py @@ -7,19 +7,23 @@ class TestAlgo(TaskPerceiver): def __init__(self): super().__init__(canny_low=((0, 255), 100), canny_high=((0, 255), 200)) + self.t = .1 def analyze(self, frame: np.ndarray, debug: bool, slider_vals: Dict[str, int]): fig = plt.figure() - x1 = np.linspace(0.0, 5.0) - x2 = np.linspace(0.0, 2.0) - - y1 = np.cos(2 * np.pi * x1) * np.exp(-x1) - y2 = np.cos(2 * np.pi * x2) - - line1, = plt.plot(x1, y1, 'ko-') - line1.set_ydata(np.cos(2 * np.pi * (x1 + slider_vals['canny_low'] * 3.14 / 2)) * np.exp(-x1)) + x = np.linspace(0.0, 5.0) + y = np.cos(2 * np.pi * (x + slider_vals['canny_low'] * 3.14 / 2)) * np.exp(-x * self.t) + plt.plot(x, y, 'ko-') fig.canvas.draw() - return frame, [frame, cv.cvtColor(frame, cv.COLOR_BGR2GRAY), cv.flip(cv.cvtColor(frame, cv.COLOR_BGR2GRAY), cv.ROTATE_180), + self.t *= 1.01 + return frame, [frame, + cv.cvtColor(frame, cv.COLOR_BGR2GRAY), + cv.flip(cv.cvtColor(frame, cv.COLOR_BGR2GRAY), cv.ROTATE_180), cv.Canny(frame, slider_vals['canny_low'], slider_vals['canny_high']), - cv.flip(cv.Canny(frame, slider_vals['canny_low'], slider_vals['canny_high']), 0), fig] \ No newline at end of file + cv.flip(cv.Canny(frame, slider_vals['canny_low'], slider_vals['canny_high']), 0), + fig] + +if __name__ == '__main__': + from perception.vis.vis import run + run(['webcam'], TestAlgo(), True) \ No newline at end of file diff --git a/perception/vis/window_builder.py b/perception/vis/Visualizer.py similarity index 86% rename from perception/vis/window_builder.py rename to perception/vis/Visualizer.py index 664db0e..65c7bb2 100644 --- a/perception/vis/window_builder.py +++ b/perception/vis/Visualizer.py @@ -2,6 +2,7 @@ import cv2 as cv import math from typing import Dict, Tuple, List +from matplotlib.pyplot import Figure def nothing(x): pass @@ -40,6 +41,15 @@ def reshape(self, frames: List[np.ndarray]) -> List[np.ndarray]: def display(self, frames: List[np.ndarray]) -> np.ndarray: num_frames = len(frames) assert (num_frames > 0 and num_frames <= 9), 'Invalid number of frames!' + + for i, frame in enumerate(frames): + if isinstance(frame, Figure): + img = np.fromstring(frame.canvas.tostring_rgb(), dtype=np.uint8, + sep='') + img = img.reshape(frame.canvas.get_width_height()[::-1] + (3,)) + img = cv.cvtColor(img, cv.COLOR_RGB2BGR) + frames[i] = img + frames = self.reshape(self.three_stack(frames)) columns = math.ceil(num_frames/math.sqrt(num_frames)) diff --git a/perception/vis/algo_stats b/perception/vis/algo_stats deleted file mode 100644 index 1ce1d89..0000000 Binary files a/perception/vis/algo_stats and /dev/null differ diff --git a/perception/vis/vis.py b/perception/vis/vis.py index ceab0f0..ee6a12b 100644 --- a/perception/vis/vis.py +++ b/perception/vis/vis.py @@ -4,84 +4,83 @@ from perception import ALGOS from perception.vis.FrameWrapper import FrameWrapper import cv2 as cv -from perception.vis.window_builder import Visualizer -import cProfile as cp -import pstats +from perception.vis.Visualizer import Visualizer +import cProfile import imageio -from matplotlib.pyplot import Figure -import numpy as np -# Parse arguments -parser = argparse.ArgumentParser(description='Visualizes perception algorithms.') -parser.add_argument('--data', default='webcam', type=str) -parser.add_argument('--algorithm', type=str, required=True) -parser.add_argument('--cProfiler', default='disabled_cprof', type=str) -parser.add_argument('--save_video', action='store_true') -args = parser.parse_args() -# Get algorithm module -Algorithm = ALGOS[args.algorithm] +def run(data_sources, algorithm, save_video=False): + out = None + window_builder = Visualizer(algorithm.kwargs) + data = FrameWrapper(data_sources, 0.25) + frame_count = 0 + speed = 1 -# Initialize image source -# detects args.data, get a list of all file directory when given a directory -# change data_source to a list of all files in the directory -if os.path.isdir(args.data): - data_sources = os.listdir(args.data) -else: - data_sources = [args.data] -data = FrameWrapper(data_sources, 0.25) + for frame in data: + if frame_count % speed == 0: + if algorithm.kwargs: + state, debug_frames = algorithm.analyze(frame, debug=True, slider_vals=window_builder.update_vars()) + else: + state, debug_frames = algorithm.analyze(frame, debug=True) -algorithm = Algorithm() -window_builder = Visualizer(algorithm.var_info()) -video_frames = [] + to_show = window_builder.display(debug_frames) + cv.imshow('Debug Frames', to_show) + if save_video: + if out is None: + out = imageio.get_writer('vis_rec.mp4') + out_img = cv.cvtColor(to_show, cv.COLOR_BGR2RGB) + out.append_data(out_img) + frame_count += 1 + key = cv.waitKey(30) + if key == ord('q') or key == 27: + break + if key == ord('p'): + cv.waitKey(0) # pause + # TODO: be able to quit and manipulate slider vars in real time while paused + if key == ord('i') and speed > 1: + speed -= 1 + print(f'speed {speed}') + if key == ord('o'): + speed += 1 + print(f'speed {speed}') -# Main Loop -def main(): - for frame in data: - state, debug_frames = algorithm.analyze( - frame, debug=True, slider_vals=window_builder.update_vars() - ) + cv.destroyAllWindows() + if out: + out.close() - for i, dframe in enumerate(debug_frames): - if isinstance(dframe, Figure): - img = np.fromstring(dframe.canvas.tostring_rgb(), dtype=np.uint8, - sep='') - img = img.reshape(dframe.canvas.get_width_height()[::-1] + (3,)) - img = cv.cvtColor(img, cv.COLOR_RGB2BGR) - debug_frames[i] = img - to_show = window_builder.display(debug_frames) - cv.imshow('Debug Frames', to_show) - if args.save_video: - video_frames.append(to_show) +def profile(*args, stats='all'): + with cProfile.Profile() as pr: + run(*args) + if stats == 'all': + pr.print_stats() + else: + pr.print_stats(stats) - key_pressed = cv.waitKey(60) & 0xFF - if key_pressed == 112: - cv.waitKey(0) # pause - if key_pressed == 113: - break # quit -if args.cProfiler == 'disabled_cprof': - main() -else: - cp.run('main()', 'algo_stats') - p = pstats.Stats('algo_stats') - if args.cProfiler != 'all_methods': - p.print_stats(args.cProfiler) - else: - p.print_stats() +if __name__ == '__main__': + # Parse arguments + parser = argparse.ArgumentParser(description='Visualizes perception algorithms.') + parser.add_argument('--data', default='webcam', type=str) + parser.add_argument('--algorithm', type=str, required=True) + parser.add_argument('--profile', default=None, type=str) + parser.add_argument('--save_video', action='store_true') + args = parser.parse_args() + + # Get algorithm class and init + algorithm = ALGOS[args.algorithm]() -cv.destroyAllWindows() + # Initialize image source + # detects args.data, get a list of all file directory when given a directory + # change data_source to a list of all files in the directory + if os.path.isdir(args.data): + data_sources = os.listdir(args.data) + else: + data_sources = [args.data] -if args.save_video: - height, width, _ = video_frames[0].shape - w = imageio.get_writer('deb_cap.mp4') - for img in video_frames: - height2, width2, _ = img.shape - if (height2, width2) == (height, width): - imag = cv.resize(img, (width - (width % 16), height - (height % 16))) - imag = cv.cvtColor(imag, cv.COLOR_BGR2RGB) - w.append_data(imag) - w.close() + if args.profile is None: + run(data_sources, algorithm, args.save_video) + else: + profile(data_sources, algorithm, args.save_video, stats=args.profile) diff --git a/requirements.txt b/requirements.txt index 6531dbb..f411292 100755 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ -opencv-python -scipy -numpy -matplotlib -imageio +opencv-python +scipy +numpy +matplotlib +imageio +imageio-ffmpeg +pytest \ No newline at end of file diff --git a/wiki/flowchart.png b/wiki/flowchart.png deleted file mode 100644 index 1bae9db..0000000 Binary files a/wiki/flowchart.png and /dev/null differ