108 lines
3.0 KiB
Python
108 lines
3.0 KiB
Python
import pyrealsense2.pyrealsense2 as rs
|
|
import numpy as np
|
|
import cv2
|
|
from threading import Thread
|
|
import logging
|
|
import time
|
|
|
|
|
|
logging.basicConfig(level='INFO')
|
|
|
|
|
|
class DepthStream:
|
|
def __init__(self) -> None:
|
|
self.depth_stream = self.init_depth_stream()
|
|
self.start_depth_stream()
|
|
self.buffer = self.get_new_depth_image()
|
|
self.updating = False
|
|
self.new_frame_ready = False
|
|
|
|
def init_depth_stream(self):
|
|
self.pipeline = rs.pipeline()
|
|
self.config = rs.config()
|
|
self.config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
|
|
|
|
def start_depth_stream(self):
|
|
self.pipeline.start(self.config)
|
|
|
|
def get_new_depth_image(self):
|
|
frames = self.pipeline.wait_for_frames()
|
|
depth_frame = frames.get_depth_frame()
|
|
depth_image = np.asanyarray(depth_frame.get_data())
|
|
return depth_image
|
|
|
|
def update(self):
|
|
while self.updating:
|
|
new_frame = self.get_new_depth_image()
|
|
self.buffer = new_frame
|
|
self.new_frame_ready = True
|
|
|
|
def start(self):
|
|
t = Thread(target=self.update, args=())
|
|
t.daemon = True
|
|
self.updating = True
|
|
t.start()
|
|
self.start_time = time.time()
|
|
return self
|
|
|
|
def read(self):
|
|
while not self.new_frame_ready:
|
|
pass
|
|
self.new_frame_ready = False
|
|
return self.buffer
|
|
|
|
def stop(self):
|
|
self.updating = False
|
|
self.pipeline.stop()
|
|
|
|
|
|
class HoleDetector:
|
|
def __init__(self) -> None:
|
|
self.thresh_value = 75
|
|
self.thresh_delta = 5
|
|
self.current_contours = []
|
|
self.running = False
|
|
self.alpha = 0.1
|
|
|
|
def start(self):
|
|
self.running = True
|
|
self.stream = DepthStream().start()
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
self.stream.stop()
|
|
|
|
def get_new_frame(self):
|
|
raw_frame = self.stream.read()
|
|
return raw_frame
|
|
|
|
def preprocess_frame(self, raw_frame):
|
|
gray_colormap = cv2.convertScaleAbs(raw_frame, alpha=self.alpha)
|
|
# NOTE: night still need to do a bgr2gray to minmax the frame
|
|
blur = cv2.medianBlur(gray_colormap, 21)
|
|
_, thresh = cv2.threshold(blur, 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)
|
|
cv2.imshow("gray_colormap", gray_colormap)
|
|
cv2.imshow("blur", blur)
|
|
cv2.imshow("thresh", thresh)
|
|
cv2.imshow("canny", canny)
|
|
return canny
|
|
|
|
|
|
def find_holes(self):
|
|
pass
|
|
|
|
def run(self):
|
|
self.start()
|
|
try:
|
|
while self.running:
|
|
new_frame = self.get_new_frame()
|
|
preprocessed_frame = self.preprocess_frame(new_frame)
|
|
cv2.waitKey(1)
|
|
self.stop()
|
|
finally:
|
|
self.stop()
|
|
|
|
if __name__ == "__main__":
|
|
a = HoleDetector().run()
|