import cv2 import numpy as np import time # stream = cv2.VideoCapture("/Users/max/Desktop/THEA_CODE/Control/Jetson/tests/depth_vid.mp4") stream = cv2.VideoCapture("/Users/max/Desktop/THEA_CODE/Control/Jetson/tests/pallet_depth.mov") # ^^ somehow we need to pipe depth data into a minmaxed greyscale frame... THRESH_VALUE = 75 class HoleFinder: def __init__(self) -> None: self.thresh_value = 75 self.thresh_delta = 5 self.current_contours = [] self.stream = cv2.VideoCapture("/Users/max/Desktop/THEA_CODE/Control/Jetson/tests/pallet_depth.mov") def get_next_frame(self): _, big_frame = self.stream.read() frame = cv2.resize(big_frame, (852, 480)) # resize needs to be appropriate for speed optimizations return frame def disp_img(self, name, frame): cv2.imshow(name, frame) def grey_blur(self, frame): gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # this normalizes blur = cv2.medianBlur(gray, 21) # This is to prevent contouring issues and make them smooth return blur def thresh_canny(self, frame): _, thresh = cv2.threshold(frame, self.thresh_value, 255, cv2.THRESH_BINARY_INV) # if we are clipping, have it move the 75 up by 5 and down by 5, and save new threah value that gave 2 correct area contours canny = cv2.Canny(thresh, 75, 200) return canny def find_contours(self, canny_frame): contours, hierarchy = cv2.findContours(canny_frame, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) self.current_contours = [] for contour in contours: area = cv2.contourArea(contour) if 5000 < area < 11000: self.current_contours.append(contour) self.current_contours = self.current_contours[::2] def cosmetics(self, source_frame): msg = "Total holes: {}".format(len(self.current_contours)) cv2.putText(source_frame, msg, (20, 40), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 2, cv2.LINE_AA) cv2.drawContours(source_frame, self.current_contours, -1, (0, 255, 0), 2) for ix, i in enumerate(self.current_contours): cv2.fillPoly(source_frame, pts = [i], color=(0, 0,255)) def evaluate_centres(self, source_frame): # loop over the contours for contour in self.current_contours: # compute the center of the contour M = cv2.moments(contour) cX = int(M["m10"] / M["m00"]) cY = int(M["m01"] / M["m00"]) cv2.circle(source_frame, (cX, cY), 7, (255, 255, 255), -1) print(f"### {cX}, {cY}") cv2.line(source_frame, (cX-50, cY), (cX+50, cY), color=(255, 128, 0), thickness=3) cv2.line(source_frame, (cX, cY-50), (cX, cY+50), color=(255, 255, 0), thickness=3) def evaluate_contours(self, blur_frame): canny_frame = self.thresh_canny(blur_frame) self.disp_img("canny_orig", canny_frame) self.find_contours(canny_frame) if len(self.current_contours) < 2: self.correct_for_offset(blur_frame) def correct_for_offset(self, blur_frame): orig_thresh = self.thresh_value self.thresh_value = orig_thresh + 0.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 0.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 1*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 1*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 1.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 1.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 2*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 2*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 2.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 2.5*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 3*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 3*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh + 4*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh - 4*self.thresh_delta success = self.test_contours(blur_frame) if success: return self.evaluate_contours(blur_frame) else: self.thresh_value = orig_thresh # give up def test_contours(self, blur_frame): canny_frame = self.thresh_canny(blur_frame) self.disp_img("canny_test", canny_frame) self.find_contours(canny_frame) if len(self.current_contours) == 2: return True return False def main(self): # time.sleep(5) while True: source_frame = self.get_next_frame() blur_frame = self.grey_blur(source_frame) self.evaluate_contours(blur_frame) self.cosmetics(source_frame) self.evaluate_centres(source_frame) self.disp_img("main", source_frame) cv2.waitKey(1) a = HoleFinder() a.main() # while True: # _, big_frame = stream.read() # # frame = cv2.resize(big_frame, (852, 480)) # resize needs to be appropriate for speed optimizations # # # cv2.imshow("orig", frame) # # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # this normalizes # # blur = cv2.medianBlur(gray, 21)#.astype('uint8') # # cv2.imshow("blur", blur) # # _, thresh = cv2.threshold(blur, 75, 255, cv2.THRESH_BINARY_INV) # if we are clipping, have it move the 75 up by 5 and down by 5, and save new threah value that gave 2 correct area contours # # cv2.imshow("thresh", thresh) # # canny = cv2.Canny(thresh, 75, 200) # # cv2.imshow('canny', canny) # contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) # contour_list = [] # print("NEW_FRAME-----") # for contour in contours: # # approx = cv2.approxPolyDP(contour, 0.04* cv2.arcLength(contour, True), True) # area = cv2.contourArea(contour) # if 5000 < area < 11000: # contour_list.append(contour) # contour_list = contour_list[::2] # print(len(contour_list)) # msg = "Total holes: {}".format(len(contour_list)) # cv2.putText(frame, msg, (20, 40), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 2, cv2.LINE_AA) # cv2.drawContours(frame, contour_list, -1, (0, 255, 0), 2) # for ix, i in enumerate(contour_list): # cv2.fillPoly(frame, pts = [i], color=(0, 0,255)) # # loop over the contours # # for contour in contour_list: # # # compute the center of the contour # # M = cv2.moments(contour) # # cX = int(M["m10"] / M["m00"]) # # cY = int(M["m01"] / M["m00"]) # # cv2.circle(frame, (cX, cY), 7, (255, 255, 255), -1) # # print(f"### {cX}, {cY}") # # cv2.line(frame, (cX-50, cY), (cX+50, cY), color=(255, 128, 0), thickness=2) # # cv2.line(frame, (cX, cY-50), (cX, cY+50), color=(255, 128, 0), thickness=2) # # print() # # print(len(contour_list)) # cv2.imshow('Objects Detected', frame) # cv2.waitKey(0)