From c63d620474e5659447ff8c9bf1f1c9eec59726a6 Mon Sep 17 00:00:00 2001 From: Daniel Parra Date: Thu, 29 Feb 2024 16:12:23 -0500 Subject: [PATCH] Detection is now possible in bright and dark environments --- .../nodes/orange_pill_detection.py | 149 +++++++++++++----- 1 file changed, 109 insertions(+), 40 deletions(-) diff --git a/SubjuGator/perception/subjugator_perception/nodes/orange_pill_detection.py b/SubjuGator/perception/subjugator_perception/nodes/orange_pill_detection.py index 679405100..9de1aa42a 100755 --- a/SubjuGator/perception/subjugator_perception/nodes/orange_pill_detection.py +++ b/SubjuGator/perception/subjugator_perception/nodes/orange_pill_detection.py @@ -5,77 +5,115 @@ import numpy as np import rospy from image_geometry import PinholeCameraModel -from mil_ros_tools import ( - Image_Publisher, - Image_Subscriber, -) +from mil_ros_tools import Image_Publisher, Image_Subscriber class CannyOrangePillDetection: - ERROR = 0.01 + THRESHOLD = 120 + ERROR = 0.05 + BRIGHTNESS_THRESHOLD = 200 + BRIGHT_OFFSET = 45 def __init__(self): - camera = rospy.get_param("~image_topic", "/camera/down/image_rect_color") - + # Use camera data + camera = rospy.get_param("~image_topic", "/camera/down/image_color") self.image_sub = Image_Subscriber(camera, self.vectorize_image) self.image_pub = Image_Publisher("~vector_viz_topic") + self.image_pub_pre = Image_Publisher("~shark_viz") self.camera_info = self.image_sub.wait_for_camera_info() + assert self.camera_info is not None self.cam = PinholeCameraModel() self.cam.fromCameraInfo(self.camera_info) - def vectorize_image(self, msg): - # Create Image from array - starttime = datetime.datetime.now() - img = msg.astype("uint8") - - img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + def binarize_frame_by_red_channel(self, frame, threshold=100): + """ + Binarize the frame based on the values of the red channel. + """ + red_channel = frame[:, :, 0] + binary_mask = (red_channel >= threshold).astype(np.uint8) + binary_mask = binary_mask * 255 + return binary_mask + + def binarize_frame_by_thresholding(self, frame): + gaus_image = cv2.GaussianBlur(frame, (3, 3), 5) + blue, green, red = cv2.split(gaus_image) + mask1 = ( + (blue - self.BRIGHT_OFFSET < red) & (blue - self.BRIGHT_OFFSET < green) + ).astype(np.uint8) * 255 + kernel = np.ones((11, 11), np.uint8) + mask1 = cv2.morphologyEx(mask1, cv2.MORPH_CLOSE, kernel) + return mask1 + + def preprocess_frame(self, frame): + """ + Extracts the red channel from the frame and magnifies its values to uncover path markers. + """ + original_frame = frame + inverted_image = cv2.bitwise_not(original_frame) + + original_array = np.array(original_frame, dtype=np.int64) + inv_original_array = np.array(inverted_image, dtype=np.int64) + + red_channel = original_array[:, :, 2] + result_image = original_array.copy() + + # MIN MAX NORM + result_image[:, :, 0] = ( + (result_image[:, :, 0] - np.min(result_image[:, :, 0])) + / (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0])) + * 255 + ) - # Apply GaussianBlur to the image to reduce noise - img_rgb = cv2.GaussianBlur(img_rgb, (5, 5), 0) + red_channel = result_image[:, :, 2] - # Make Mask - mask = (img_rgb[:, :, 2] / img_rgb[:, :, 0] > 0.5) | ( - img_rgb[:, :, 2] / img_rgb[:, :, 1] > 0.5 + result_image[:, :, 0] = red_channel * ( + red_channel + inv_original_array[:, :, 0] + ) + result_image[:, :, 1] = 0 + result_image[:, :, 2] = 0 + + # MIN MAX NORM + result_image[:, :, 0] = ( + (result_image[:, :, 0] - np.min(result_image[:, :, 0])) + / (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0])) + * 255 ) - # Set pixels where the condition is true to black - img_rgb[mask] = [0, 0, 0] + result_image = np.clip(result_image, 0, 255) + result_image = result_image.astype(np.uint8) - # Apply Canny edge detection - edges = cv2.Canny(img_rgb, 50, 150) + self.image_pub_pre.publish(np.array(result_image)) - # Vector Analysis - # Find contours in the binary image + return result_image + + def extract_edges(self, bin_img): + edges = cv2.Canny(bin_img, 20, 200) contours, _ = cv2.findContours( edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE, ) - - # Create an empty image to draw contours on contour_image = np.zeros_like(edges) - contour_image_bgr = cv2.cvtColor(contour_image, cv2.COLOR_GRAY2BGR) - - # Draw contours on the empty image - cv2.drawContours(contour_image, contours, -1, (255), 2) # -1 draws all contours - + cv2.drawContours(contour_image, contours, -1, (255), 2) lines = cv2.HoughLinesP( contour_image, 1, np.pi / 180, threshold=50, - minLineLength=30, - maxLineGap=10, + minLineLength=100, + maxLineGap=20, ) + return lines, contour_image_bgr + + def group_and_draw_vectors(self, lines, contour_image_bgr): + radians = {} if lines is None: - bgr_image = cv2.cvtColor(contour_image, cv2.COLOR_GRAY2BGR) + bgr_image = contour_image_bgr self.image_pub.publish(np.array(bgr_image)) - return - - radians = {} + return radians for i in range(len(lines)): line = lines[i] @@ -113,7 +151,38 @@ def vectorize_image(self, msg): sorted_dict = dict(sorted(radians.items(), key=lambda item: -len(item[1]))) - # DIsplay vectors on image + return sorted_dict + + def vectorize_image(self, msg): + # Create Image from array + datetime.datetime.now() + img = msg.astype("uint8") + + # Calculate the average brightness + grayscale_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + average_brightness = cv2.mean(grayscale_image)[0] + + bin_frame = None + + if average_brightness > self.BRIGHTNESS_THRESHOLD: + bin_frame = self.binarize_frame_by_thresholding(frame=img) + + else: + preprocessed_frame = self.preprocess_frame(frame=img) + + bin_frame = self.binarize_frame_by_red_channel( + preprocessed_frame, + self.THRESHOLD, + ) + + lines, contour_image_bgr = self.extract_edges(bin_img=bin_frame) + + sorted_dict = self.group_and_draw_vectors( + lines=lines, + contour_image_bgr=contour_image_bgr, + ) + + # Display vectors on image for key, value in sorted_dict.items(): major_angle = sum(value) / len(value) print(major_angle) @@ -122,8 +191,8 @@ def vectorize_image(self, msg): # bgr_image = cv2.cvtColor(contour_image_bgr) self.image_pub.publish(np.array(contour_image_bgr)) - endtime = datetime.datetime.now() - print(endtime - starttime) + datetime.datetime.now() + # print(endtime - starttime) if __name__ == "__main__":