From ae6646dfa28b05ade2bcf05ea3707073de92038c Mon Sep 17 00:00:00 2001 From: David Chu Date: Wed, 1 Jun 2022 22:40:28 -0400 Subject: [PATCH 1/3] Linted with Black (sorry... it's automatic and my brain is too fried to have turned it off), added error handling (both try-except and if-else clauses), and optimized the preprocessing of lane detection a little bit. --- cv/cv/src/cv_utils/camera_projection.py | 55 ++++--- .../src/lane_detection_model_inference.py | 135 +++++++++--------- .../src/lane_detection_wrapper_client.py | 30 ++-- .../src/pothole_detection_model_inference.py | 36 +++-- .../src/pothole_detection_wrapper_client.py | 63 ++++---- 5 files changed, 169 insertions(+), 150 deletions(-) diff --git a/cv/cv/src/cv_utils/camera_projection.py b/cv/cv/src/cv_utils/camera_projection.py index e4745f72..39e6dd61 100644 --- a/cv/cv/src/cv_utils/camera_projection.py +++ b/cv/cv/src/cv_utils/camera_projection.py @@ -1,21 +1,21 @@ -import json -import os +#!/usr/bin/env python2 +import json +import os -import cv2 import numpy as np +import rospkg import rospy from image_geometry import PinholeCameraModel -import rospkg - from sensor_msgs.msg import CameraInfo + class CameraProjection: def __init__(self): rospack = rospkg.RosPack() - save_path = rospack.get_path('cv') + '/config/depth_vals.json' + save_path = rospack.get_path("cv") + "/config/depth_vals.json" if os.path.exists(save_path): print("Using depth map values from", save_path) - self.depth_map = json.load(open(save_path, 'r')) + self.depth_map = json.load(open(save_path, "r")) else: self.depth_map = self.gen_depth_to_y_map(1.5, 4) @@ -26,11 +26,13 @@ def __init__(self): self.depth_values += [self.depth_map[key]] self.depth_values = np.array(self.depth_values) - camera_info = rospy.wait_for_message('/zed/zed_node/rgb/camera_info', CameraInfo) + camera_info = rospy.wait_for_message( + "/zed/zed_node/rgb/camera_info", CameraInfo + ) # Update camera intrinsics to account for the resized new images - scale_x = 330.0/camera_info.width - scale_y = 180.0/camera_info.height + scale_x = 330.0 / camera_info.width + scale_y = 180.0 / camera_info.height camera_info.width = 330 camera_info.height = 180 @@ -45,7 +47,7 @@ def __init__(self): P = list(camera_info.P) P[0] = P[0] * scale_x P[2] = P[2] * scale_x - P[3] = P[3] * scale_x + P[3] = P[3] * scale_x P[5] = P[5] * scale_y P[6] = P[6] * scale_y P[7] = P[7] * scale_y @@ -60,44 +62,41 @@ def __init__(self): self.camera.fromCameraInfo(camera_info) - def __call__(self, pts): - ''' + """ [[x, y], [x, y]] - ''' + """ return self.project_2D_to_3D_camera(pts) - def project_2D_to_3D_camera(self, pts): - ''' Given the pixel image coordinates and a constant mapping of depth to y-values, - calculate the resulting projection of points in the camera coordinate frame. - Requires the focal lengths in x and y and the optical centers - ''' + """Given the pixel image coordinates and a constant mapping of depth to y-values, + calculate the resulting projection of points in the camera coordinate frame. + Requires the focal lengths in x and y and the optical centers + """ points = [] for i in range(len(pts)): if pts[i][0] < 0 or pts[i][1] < 0 or pts[i][0] >= 330 or pts[i][1] >= 180: continue vec = self.camera.projectPixelTo3dRay((pts[i][0], pts[i][1])) z_val = self.depth_map[str((pts[i][0], pts[i][1]))] - point3D = [vec[i] * z_val/vec[2] for i in range(3)] + point3D = [vec[i] * z_val / vec[2] for i in range(3)] points += [point3D] return points - def gen_depth_to_y_map(self, min_z, max_z): - ''' Create the constant depth to y map ''' + """Create the constant depth to y map""" depth_map = {} - delta_z = (max_z-min_z)/6 + delta_z = (max_z - min_z) / 6 for i, _y in enumerate(range(150, -30, -30)): for _x in range(0, 330, 30): - for y in range(_y, _y+30): - for x in range(_x, _x+30): - depth_map[str((x, y))] = delta_z*i + for y in range(_y, _y + 30): + for x in range(_x, _x + 30): + depth_map[str((x, y))] = delta_z * i - for n,y in enumerate(range(180-1,-1,-1),1): - depth_map[y] = delta_z*n + for n, y in enumerate(range(180 - 1, -1, -1), 1): + depth_map[y] = delta_z * n return depth_map diff --git a/cv/lane_detection/src/lane_detection_model_inference.py b/cv/lane_detection/src/lane_detection_model_inference.py index 85866053..cc95a89e 100755 --- a/cv/lane_detection/src/lane_detection_model_inference.py +++ b/cv/lane_detection/src/lane_detection_model_inference.py @@ -1,136 +1,143 @@ #!/usr/bin/env python3 -import struct -import sys -import time import json -import os +import struct import cv2 import numpy as np -import redis import onnx -import onnxruntime as ort - +import onnxruntime as ort +import redis import rospkg from line_fitting import fit_lanes + class CVModelInferencer: def __init__(self): - self.redis = redis.Redis(host='127.0.0.1', port=6379, db=3) - + self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) + rospack = rospkg.RosPack() - model_path = rospack.get_path('lane_detection') + '/models/unet_with_sigmoid.onnx' + model_path = ( + rospack.get_path("lane_detection") + "/models/unet_with_sigmoid.onnx" + ) self.model = onnx.load(model_path) onnx.checker.check_model(self.model) - self.ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider']) + self.ort_session = ort.InferenceSession( + model_path, providers=["CUDAExecutionProvider"] + ) def run(self): raw = self._fromRedisImg("zed/preprocessed") - if raw is not None: - img = get_input(raw.copy()) - - # Do model inference - output = self.ort_session.run(None, {'Inputs': img.astype(np.float32)})[0][0][0] - mask = np.where(output > 0.5, 1., 0.) - self._toRedisImg(mask, "cv/model/output") + if raw is None: + return + img = get_input(raw) # Does not return None - lanes = fit_lanes(mask) - if lanes is not None: - self._toRedisLanes(lanes, "lane_detection") + # Do model inference + outputs = self.ort_session.run(None, {"Inputs": img.astype(np.float32)}) + mask = np.where(outputs[0][0][0] > 0.5, 1.0, 0.0) + self._toRedisImg(mask, "cv/model/output") - # toshow = np.concatenate([cv2.resize(raw, (330, 180)), np.tile(mask[...,np.newaxis]*255, (1, 1, 3))], axis=1).astype(np.uint8) - # cv2.imshow("image", toshow) - # cv2.waitKey(1) + lanes = fit_lanes(mask) + if lanes is not None: + self._toRedisLanes(lanes, "lane_detection") - def _toRedisLanes(self,lanes,name): + def _toRedisLanes(self, lanes, name): """Store given Numpy array 'img' in Redis under key 'name'""" encoded = json.dumps(lanes) # Store encoded data in Redis - self.redis.set(name,encoded) + self.redis.set(name, encoded) def _fromRedisImg(self, name): """Retrieve Numpy array from Redis key 'n'""" encoded = self.redis.get(name) if encoded is None: return None - h, w = struct.unpack('>II',encoded[:8]) - a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape(h,w,3) + h, w = struct.unpack(">II", encoded[:8]) + a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape( + h, w, 3 + ) return a - def _toRedisImg(self,img,name): + def _toRedisImg(self, img, name): """Store given Numpy array 'img' in Redis under key 'name'""" h, w = img.shape[:2] - shape = struct.pack('>II',h,w) + shape = struct.pack(">II", h, w) - retval, buffer = cv2.imencode('.png', img) + retval, buffer = cv2.imencode(".png", img) img_bytes = np.array(buffer).tostring() encoded = shape + img_bytes # Store encoded data in Redis - self.redis.set(name,encoded) + self.redis.set(name, encoded) return + def find_edge_channel(img): - edges_mask = np.zeros((img.shape[0],img.shape[1]),dtype=np.uint8) + edges_mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8) width = img.shape[1] height = img.shape[0] - gray_im = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) + gray_im = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # gray_im = cv2.GaussianBlur(gray_im,(3,3),0) # Separate into quadrants - med1 = np.median(gray_im[:height//2,:width//2]) - med2 = np.median(gray_im[:height//2,width//2:]) - med3 = np.median(gray_im[height//2:,width//2:]) - med4 = np.median(gray_im[height//2:,:width//2]) + med1 = np.median(gray_im[: height // 2, : width // 2]) + med2 = np.median(gray_im[: height // 2, width // 2 :]) + med3 = np.median(gray_im[height // 2 :, width // 2 :]) + med4 = np.median(gray_im[height // 2 :, : width // 2]) - l1 = int(max(0,(1-0.205)*med1)) - u1 = int(min(255,(1+0.205)*med1)) - e1 = cv2.Canny(gray_im[:height//2,:width//2],l1,u1) + l1 = int(max(0, (1 - 0.205) * med1)) + u1 = int(min(255, (1 + 0.205) * med1)) + e1 = cv2.Canny(gray_im[: height // 2, : width // 2], l1, u1) - l2 = int(max(0,(1-0.205)*med2)) - u2 = int(min(255,(1+0.205)*med2)) - e2 = cv2.Canny(gray_im[:height//2,width//2:],l2,u2) + l2 = int(max(0, (1 - 0.205) * med2)) + u2 = int(min(255, (1 + 0.205) * med2)) + e2 = cv2.Canny(gray_im[: height // 2, width // 2 :], l2, u2) - l3 = int(max(0,(1-0.205)*med3)) - u3 = int(min(255,(1+0.205)*med3)) - e3 = cv2.Canny(gray_im[height//2:,width//2:],l3,u3) + l3 = int(max(0, (1 - 0.205) * med3)) + u3 = int(min(255, (1 + 0.205) * med3)) + e3 = cv2.Canny(gray_im[height // 2 :, width // 2 :], l3, u3) - l4 = int(max(0,(1-0.205)*med4)) - u4 = int(min(255,(1+0.205)*med4)) - e4 = cv2.Canny(gray_im[height//2:,:width//2],l4,u4) + l4 = int(max(0, (1 - 0.205) * med4)) + u4 = int(min(255, (1 + 0.205) * med4)) + e4 = cv2.Canny(gray_im[height // 2 :, : width // 2], l4, u4) # Stitch the edges together - edges_mask[:height//2,:width//2] = e1 - edges_mask[:height//2,width//2:] = e2 - edges_mask[height//2:,width//2:] = e3 - edges_mask[height//2:,:width//2] = e4 + edges_mask[: height // 2, : width // 2] = e1 + edges_mask[: height // 2, width // 2 :] = e2 + edges_mask[height // 2 :, width // 2 :] = e3 + edges_mask[height // 2 :, : width // 2] = e4 edges_mask_inv = cv2.bitwise_not(edges_mask) return edges_mask, edges_mask_inv -def get_input(frame): - frame = cv2.resize(frame,(1280,720),interpolation=cv2.INTER_AREA) +def get_input(frame: np.ndarray) -> np.ndarray: + frame = cv2.resize(frame, (330, 180), interpolation=cv2.INTER_AREA) frame_copy = np.copy(frame) - test_edges,test_edges_inv = find_edge_channel(frame_copy) - frame_copy = np.append(frame_copy,test_edges.reshape(test_edges.shape[0],test_edges.shape[1],1),axis=2) - frame_copy = np.append(frame_copy,test_edges_inv.reshape(test_edges_inv.shape[0],test_edges_inv.shape[1],1),axis=2) - frame_copy = cv2.resize(frame_copy,(330,180)) + test_edges, test_edges_inv = find_edge_channel(frame_copy) + frame_copy = np.block( + [ + frame_copy, + test_edges.reshape(test_edges.shape[0], test_edges.shape[1], 1), + test_edges_inv.reshape(test_edges_inv.shape[0], test_edges_inv.shape[1], 1), + ] + ) - input = (frame_copy/255.).transpose(2,0,1).reshape(1,5,180,330) + input = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 5, 180, 330) return input -if __name__ == '__main__': +if __name__ == "__main__": wrapper = CVModelInferencer() - while True: - wrapper.run() \ No newline at end of file + try: + wrapper.run() + except: # Top-level loop must not fail + pass diff --git a/cv/lane_detection/src/lane_detection_wrapper_client.py b/cv/lane_detection/src/lane_detection_wrapper_client.py index 7a8f1234..80f98c54 100755 --- a/cv/lane_detection/src/lane_detection_wrapper_client.py +++ b/cv/lane_detection/src/lane_detection_wrapper_client.py @@ -1,6 +1,5 @@ #!/usr/bin/env python2 import struct -import sys import json import numpy as np @@ -14,24 +13,25 @@ from geometry_msgs.msg import Point from cv_utils import camera_projection + class CVWrapperClient: def __init__(self): - self.redis = redis.Redis(host='127.0.0.1', port=6379, db=3) - self.pub = rospy.Publisher('cv/lane_detections', FloatArray, queue_size=10) - self.pub_raw = rospy.Publisher('cv/model_output', Image, queue_size=10) - rospy.init_node('lane_detection_wrapper_client') - self.r = rospy.Rate(10) # 10hz + self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) + self.pub = rospy.Publisher("cv/lane_detections", FloatArray, queue_size=10) + self.pub_raw = rospy.Publisher("cv/model_output", Image, queue_size=10) + rospy.init_node("lane_detection_wrapper_client") + self.r = rospy.Rate(10) # 10hz self.bridge = CvBridge() self.projection = camera_projection.CameraProjection() def run(self): while not rospy.is_shutdown(): - lanes = self._fromRedisLanes('lane_detection') - mask = self._fromRedisImg('cv/model/output') + lanes = self._fromRedisLanes("lane_detection") + mask = self._fromRedisImg("cv/model/output") if img is not None: - img = self.bridge.cv2_to_imgmsg(mask*255, encoding='passthrough') + img = self.bridge.cv2_to_imgmsg(mask * 255, encoding="passthrough") self.pub_raw.publish(img) if lanes is not None: @@ -69,13 +69,13 @@ def _fromRedisImg(self, name): encoded = self.redis.get(name) if encoded is None: return None - h, w = struct.unpack('>II',encoded[:8]) - a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape(h,w,3) + h, w = struct.unpack(">II", encoded[:8]) + a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape( + h, w, 3 + ) return a -if __name__ == '__main__': + +if __name__ == "__main__": wrapper = CVWrapperClient() wrapper.run() - - - \ No newline at end of file diff --git a/cv/pothole_detection/src/pothole_detection_model_inference.py b/cv/pothole_detection/src/pothole_detection_model_inference.py index b992b88e..4d00d516 100755 --- a/cv/pothole_detection/src/pothole_detection_model_inference.py +++ b/cv/pothole_detection/src/pothole_detection_model_inference.py @@ -3,24 +3,26 @@ import struct import cv2 -from cv2 import intersectConvexConvex import numpy as np import onnx import onnxruntime as ort import redis import rospkg + from runonnx import run_inference class CVModelInferencer: def __init__(self): - self.redis = redis.Redis(host='127.0.0.1', port=6379, db=3) - + self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) + rospack = rospkg.RosPack() model_path = f"{rospack.get_path('pothole_detection')}/models/best.onnx" self.model = onnx.load(model_path) onnx.checker.check_model(self.model) - self.ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider']) + self.ort_session = ort.InferenceSession( + model_path, providers=["CUDAExecutionProvider"] + ) def run(self) -> None: raw = self._fromRedis("zed/preprocessed") @@ -33,7 +35,12 @@ def run(self) -> None: potholes = potholes[0] potholes = [potholes[i][:4] for i in range(len(potholes))] - potholes = [[[int(p[0] + ((p[2] - p[0]) / 2)), int(p[1] + ((p[3] - p[1]) / 2))] for p in potholes]] + potholes = [ + [ + [int(p[0] + ((p[2] - p[0]) / 2)), int(p[1] + ((p[3] - p[1]) / 2))] + for p in potholes + ] + ] # json_dumpable = {"pothole": potholes} # raw = cv2.resize(raw, (330, 180)) @@ -48,25 +55,30 @@ def _toRedis(self, lanes, name): """Store given Numpy array 'img' in Redis under key 'name'""" encoded = json.dumps(lanes) # Store encoded data in Redis - self.redis.set(name,encoded) + self.redis.set(name, encoded) def _fromRedis(self, name): """Retrieve Numpy array from Redis key 'n'""" encoded = self.redis.get(name) - h, w = struct.unpack('>II',encoded[:8]) - a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape(h,w,3) + h, w = struct.unpack(">II", encoded[:8]) + a = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8, offset=8), 1).reshape( + h, w, 3 + ) return a def get_copied_resized_input(frame: np.ndarray) -> np.ndarray: frame_copy = np.copy(frame) - frame_copy = cv2.resize(frame_copy,(448,448)) - frame_copy = (frame_copy / 255.).transpose(2,0,1).reshape(1,3,448,448) + frame_copy = cv2.resize(frame_copy, (448, 448)) + frame_copy = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 3, 448, 448) return frame_copy -if __name__ == '__main__': +if __name__ == "__main__": wrapper = CVModelInferencer() while True: - wrapper.run() + try: + wrapper.run() + except: # Top-level loop must not fail + pass diff --git a/cv/pothole_detection/src/pothole_detection_wrapper_client.py b/cv/pothole_detection/src/pothole_detection_wrapper_client.py index 683d35c2..6cb61199 100755 --- a/cv/pothole_detection/src/pothole_detection_wrapper_client.py +++ b/cv/pothole_detection/src/pothole_detection_wrapper_client.py @@ -1,56 +1,57 @@ #!/usr/bin/env python2 -import struct -import sys import json import numpy as np import redis import rospy - from cv.msg import FloatArray, FloatList -from geometry_msgs.msg import Point from cv_utils import camera_projection +from geometry_msgs.msg import Point + class CVWrapperClient: def __init__(self): - self.redis = redis.Redis(host='127.0.0.1', port=6379, db=3) - self.pub = rospy.Publisher('cv/pothole_detections', FloatArray, queue_size=10) - rospy.init_node('pothole_detection_wrapper_client') - self.r = rospy.Rate(10) # 10hz - + self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) + self.pub = rospy.Publisher("cv/pothole_detections", FloatArray, queue_size=10) + rospy.init_node("pothole_detection_wrapper_client") + self.r = rospy.Rate(10) # 10hz self.projection = camera_projection.CameraProjection() def run(self): while not rospy.is_shutdown(): - lanes = self._fromRedis('pothole_detection') - - lanes_msgs = [] - for lane in lanes: - project_points = self.projection(np.array(lane, dtype=np.int)) - - lane_msg = FloatList() - pts_msg = [] - for pt in project_points: - pt_msg = Point() - pt_msg.x = pt[0] - pt_msg.y = pt[1] - pt_msg.z = pt[2] - pts_msg += [pt_msg] - lane_msg.elements = pts_msg - lanes_msgs += [lane_msg] - - msg = FloatArray() - msg.lists = lanes_msgs - self.pub.publish(msg) - + lanes = self._fromRedis("pothole_detection") + if lanes is not None: + lanes_msgs = [] + for lane in lanes: + project_points = self.projection(np.array(lane, dtype=np.int)) + if project_points is None: + continue + + lane_msg = FloatList() # Failure here is fatal + pts_msg = [] + for pt in project_points: + pt_msg = Point() + pt_msg.x = pt[0] + pt_msg.y = pt[1] + pt_msg.z = pt[2] + pts_msg += [pt_msg] + lane_msg.elements = pts_msg + lanes_msgs += [lane_msg] + + msg = FloatArray() # Failure here is fatal + msg.lists = lanes_msgs + self.pub.publish(msg) self.r.sleep() def _fromRedis(self, name): """Retrieve Numpy array from Redis key 'n'""" encoded = self.redis.get(name) + if encoded is None: + return None lanes = json.loads(encoded) return lanes -if __name__ == '__main__': + +if __name__ == "__main__": wrapper = CVWrapperClient() wrapper.run() From 79bb99bbfab829284595dfe8d59d34cb69921257 Mon Sep 17 00:00:00 2001 From: David Chu Date: Wed, 1 Jun 2022 23:19:08 -0400 Subject: [PATCH 2/3] Rearranged the copying of the frame to prevent changing it. I also moved numbers around to do stuff --- .../src/lane_detection_model_inference.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/cv/lane_detection/src/lane_detection_model_inference.py b/cv/lane_detection/src/lane_detection_model_inference.py index cc95a89e..304493d5 100755 --- a/cv/lane_detection/src/lane_detection_model_inference.py +++ b/cv/lane_detection/src/lane_detection_model_inference.py @@ -12,6 +12,10 @@ from line_fitting import fit_lanes +WIDTH: int = 330 +HEIGHT: int = 180 + + class CVModelInferencer: def __init__(self): self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) @@ -78,7 +82,7 @@ def _toRedisImg(self, img, name): def find_edge_channel(img): - edges_mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8) + edges_mask = np.zeros_like(img, dtype=np.uint8) width = img.shape[1] height = img.shape[0] @@ -118,19 +122,21 @@ def find_edge_channel(img): def get_input(frame: np.ndarray) -> np.ndarray: - frame = cv2.resize(frame, (330, 180), interpolation=cv2.INTER_AREA) frame_copy = np.copy(frame) + frame_copy = cv2.resize(frame_copy, (WIDTH, HEIGHT), interpolation=cv2.INTER_AREA) test_edges, test_edges_inv = find_edge_channel(frame_copy) frame_copy = np.block( [ - frame_copy, - test_edges.reshape(test_edges.shape[0], test_edges.shape[1], 1), - test_edges_inv.reshape(test_edges_inv.shape[0], test_edges_inv.shape[1], 1), + frame_copy, # shape: (330, 180, 3) + test_edges.reshape(WIDTH, HEIGHT, 1), # shape: (330, 180, 1) + test_edges_inv.reshape(WIDTH, HEIGHT, 1), # shape: (330, 180, 1) ] - ) + ) # (330, 180, 5) - input = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 5, 180, 330) + # Is this the correct permutation? Maybe try it out + # Shapes go: (330, 180, 5) ----> (5, 330, 180) --ERROR??--> (1, 5, 180, 330) + input = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 5, HEIGHT, WIDTH) return input From e3053d8c42864d46271288934e7e780d7a0d665c Mon Sep 17 00:00:00 2001 From: Spencer Teetaert Date: Thu, 2 Jun 2022 00:43:59 -0400 Subject: [PATCH 3/3] Fixes + added debug mode --- .../src/lane_detection_model_inference.py | 30 ++++++++++++------- .../src/lane_detection_wrapper_client.py | 2 +- .../src/pothole_detection_model_inference.py | 28 ++++++++++------- 3 files changed, 37 insertions(+), 23 deletions(-) diff --git a/cv/lane_detection/src/lane_detection_model_inference.py b/cv/lane_detection/src/lane_detection_model_inference.py index 304493d5..9d4c2bae 100755 --- a/cv/lane_detection/src/lane_detection_model_inference.py +++ b/cv/lane_detection/src/lane_detection_model_inference.py @@ -14,6 +14,7 @@ WIDTH: int = 330 HEIGHT: int = 180 +DEBUG_MODE: bool = True class CVModelInferencer: @@ -47,6 +48,11 @@ def run(self): if lanes is not None: self._toRedisLanes(lanes, "lane_detection") + if DEBUG_MODE: + toshow = np.concatenate([cv2.resize(raw, (330, 180)), np.tile(mask[...,np.newaxis]*255, (1, 1, 3))], axis=1).astype(np.uint8) + cv2.imshow("image", toshow) + cv2.waitKey(1) + def _toRedisLanes(self, lanes, name): """Store given Numpy array 'img' in Redis under key 'name'""" @@ -82,7 +88,7 @@ def _toRedisImg(self, img, name): def find_edge_channel(img): - edges_mask = np.zeros_like(img, dtype=np.uint8) + edges_mask = np.zeros(img.shape[0:2], dtype=np.uint8) width = img.shape[1] height = img.shape[0] @@ -128,22 +134,24 @@ def get_input(frame: np.ndarray) -> np.ndarray: test_edges, test_edges_inv = find_edge_channel(frame_copy) frame_copy = np.block( [ - frame_copy, # shape: (330, 180, 3) - test_edges.reshape(WIDTH, HEIGHT, 1), # shape: (330, 180, 1) - test_edges_inv.reshape(WIDTH, HEIGHT, 1), # shape: (330, 180, 1) + frame_copy, # shape: (180, 330, 3) + test_edges.reshape(*test_edges.shape, 1), # shape: (180, 330, 1) + test_edges_inv.reshape(*test_edges_inv.shape, 1), # shape: (180, 330, 1) ] ) # (330, 180, 5) - # Is this the correct permutation? Maybe try it out - # Shapes go: (330, 180, 5) ----> (5, 330, 180) --ERROR??--> (1, 5, 180, 330) - input = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 5, HEIGHT, WIDTH) + input = (frame_copy / 255.0).transpose(2, 0, 1).reshape(1, 5, *test_edges.shape) return input if __name__ == "__main__": wrapper = CVModelInferencer() - while True: - try: + if DEBUG_MODE: + while True: wrapper.run() - except: # Top-level loop must not fail - pass + else: + while True: + try: + wrapper.run() + except: + pass diff --git a/cv/lane_detection/src/lane_detection_wrapper_client.py b/cv/lane_detection/src/lane_detection_wrapper_client.py index 80f98c54..4c477991 100755 --- a/cv/lane_detection/src/lane_detection_wrapper_client.py +++ b/cv/lane_detection/src/lane_detection_wrapper_client.py @@ -30,7 +30,7 @@ def run(self): lanes = self._fromRedisLanes("lane_detection") mask = self._fromRedisImg("cv/model/output") - if img is not None: + if mask is not None: img = self.bridge.cv2_to_imgmsg(mask * 255, encoding="passthrough") self.pub_raw.publish(img) diff --git a/cv/pothole_detection/src/pothole_detection_model_inference.py b/cv/pothole_detection/src/pothole_detection_model_inference.py index 4d00d516..d2f84bc8 100755 --- a/cv/pothole_detection/src/pothole_detection_model_inference.py +++ b/cv/pothole_detection/src/pothole_detection_model_inference.py @@ -12,6 +12,9 @@ from runonnx import run_inference +DEBUG_MODE: bool = True + + class CVModelInferencer: def __init__(self): self.redis = redis.Redis(host="127.0.0.1", port=6379, db=3) @@ -41,13 +44,13 @@ def run(self) -> None: for p in potholes ] ] - # json_dumpable = {"pothole": potholes} - # raw = cv2.resize(raw, (330, 180)) - # for circle in potholes[0]: - # raw = cv2.circle(raw, circle, 5, (0, 0, 255), -1) - # cv2.imshow("im", raw) - # cv2.waitKey(1) + if DEBUG_MODE: + raw = cv2.resize(raw, (330, 180)) + for circle in potholes[0]: + raw = cv2.circle(raw, circle, 5, (0, 0, 255), -1) + cv2.imshow("im", raw) + cv2.waitKey(1) self._toRedis(potholes, "pothole_detection") @@ -76,9 +79,12 @@ def get_copied_resized_input(frame: np.ndarray) -> np.ndarray: if __name__ == "__main__": wrapper = CVModelInferencer() - - while True: - try: + if DEBUG_MODE: + while True: wrapper.run() - except: # Top-level loop must not fail - pass + else: + while True: + try: + wrapper.run() + except: + pass