Change camera topic

SubjuGator/perception/subjugator_perception/nodes/orange_pill_detection.py

@@ -5,77 +5,114 @@
 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/front/right/image_rect_color")
-
+        camera = rospy.get_param("~image_topic", "/camera/down/image_rect_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 +150,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 +190,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__":