231 lines
10 KiB
Python
231 lines
10 KiB
Python
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)
|