Test Script for Vision Stack

test/topics_test.py

@@ -1,7 +1,8 @@
 import rclpy
 import pytest
-from PIL import Image
+
 import numpy as np
+from PIL import Image
 
 from vision_stack import VisionStack
 from vision_stack import BinThresholdingLayer, CannyLayer, ColorMagnificationLayer, CustomLayer, GaussianLayer, GrayscaleLayer, HoughTransformLayer, MinMaxNormalizationLayer, ZScoreNormalizationLayer, ObjectDetectionLayer, ResizeLayer, RGBMagnificationLayer, RGBtoBGRLayer, SobelLayer, UnderWaterImageEnhancementLayer 
@@ -30,6 +31,8 @@ def custom_process(img, args):
         SobelLayer(5)
     ], 'test')
 
+    print("Initialization is complete")
+
     img = np.array(Image.open('test/test.jpg'))
     vs.run(img, True)
     print("here2")
@@ -62,4 +65,4 @@ def custom_process(img, args):
     rclpy.shutdown()
 
 
-
+test_topic_creation()

test/vision_stack_test.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+import rclpy
+from rclpy.node import Node
+
+from sensor_msgs.msg import Image
+
+from cv_bridge import CvBridge
+
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent.parent))
+vision_stack_base_path = str(Path(__file__).parent.parent)
+
+from vision_stack import VisionStack, BinThresholdingLayer, CannyLayer, ColorMagnificationLayer, CustomLayer, GaussianLayer, GrayscaleLayer, HoughTransformLayer, MinMaxNormalizationLayer, ZScoreNormalizationLayer, RobustScalingLayer, ObjectDetectionLayer, ResizeLayer, RGBMagnificationLayer, SobelLayer, UnderWaterImageEnhancementLayer
+
+bridge = CvBridge()
+
+class vision_stack_test(Node):
+
+    def __init__(self):
+        super().__init__('Vision_Stack_Subscriber')
+
+        # Creating the subscriber to the front_cam image data
+        self.subscription = self.create_subscription(
+            Image,
+            '/front_cam/image_raw', # You can replace this with down_cam/image_raw if you want
+            self.listener_callback,
+            10)
+
+        self.vs = VisionStack(
+            layers=[
+                # Include as many layers in any combination as you need
+
+                # Stole the comments and layer initializations below from Daniel's README on the VisionStack repo
+
+                ObjectDetectionLayer(conf_thres=0.5, weights_file=f"{vision_stack_base_path}/vision_stack/weights/test.pt", pass_post_detection_img = False), # Access YOLO weights and make predictions on the image provided by the previous layer. Setting pass_post_processing_img will push the image with bounding boxes to the next layer.
+                GrayscaleLayer(), # Convert image to grayscale
+                BinThresholdingLayer(150,250), # Converts image to grayscale if image is not grayscale and extracts pixels with values between 150 and 250.
+                CannyLayer(50,100), # Simplified canny filter that uses cv2.Canny to pass a canny filter over an image with the low value (50) threshold for soft edge detection and the high value (100) for strong edges detection.
+                ColorMagnificationLayer((23,156,234)), # Highlights objects with this color (23,156,234) in an image.
+                GaussianLayer((11,11), 50), # Pass a gaussian filter of kernal size (11,11) (kernal size must be of odd numbered dimensions) with a sigma value of 50.
+                HoughTransformLayer(threshold=100, min_line_length=20, max_line_gap=10, pass_post_processing_img=True), # Pass a Hough Transform filter over an image to extract lines. Setting pass_post_processing_img will push the image with hough transform lines to the next layer.
+                MinMaxNormalizationLayer(),
+                ZScoreNormalizationLayer(),
+                ResizeLayer(960, 608), # Resize the image from the previous layer.
+                RGBMagnificationLayer('G'), # Magnifies the provided channel respectively (options are 'R', 'G', 'B')
+                UnderWaterImageEnhancementLayer(), # Uses a generative AI model to improve underwater images (good for murky waters).
+
+
+                # **********Layers with issues****************
+                # RobustScalingLayer(), This layer produces a floating point encoding type which is not allowed by the publisher, so don't use this alyer
+                # SobelLayer(5), @ Fails for the same reason as RobustScalingLayer
+            ],
+        )
+
+
+    # Runs vision stack on every callback
+    def listener_callback(self, msg):
+
+        cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
+        self.vs.run(cv_image, True)
+
+
+def main(args=None):
+    rclpy.init(args=args)
+    front_cam_subscriber = vision_stack_test()
+    rclpy.spin(front_cam_subscriber)
+
+
+
+if __name__ == "__main__":
+    main()

vision_stack/VisionStack.py

@@ -3,16 +3,18 @@
 import time
 import numpy as np
 
+import copy
+
 from cv_bridge import CvBridge, CvBridgeError
 from sensor_msgs.msg import Image
 
 try:
     import rclpy
     from rclpy.publisher import Publisher
     from rclpy.node import Node
-except:
     import matplotlib.pyplot as plt
     plt.switch_backend('TkAgg')
+except:
     print("mil_ros_tools package is not available")
 
 NUM_COLS = 3
@@ -32,6 +34,7 @@ class Image_Publisher:
     """
 
     def __init__(self, topic: str, node:Node, encoding: str = "infer", queue_size: int = 1):
+        print(f"Topic Name: {topic}")
         self.bridge = CvBridge()
         self.encoding = encoding
         self.node:Node = node
@@ -125,7 +128,7 @@ def pop(self, index = -1):
         self.layers.pop(index)
 
     def run(self, in_image, verbose = False):
-        processed_image = in_image.copy()
+        processed_image = copy.copy(in_image)
         self.analysis_dict["updated_at"] = time.localtime()
 
         num_rows = -(-len(self.layers) // NUM_COLS)
@@ -135,7 +138,7 @@ def run(self, in_image, verbose = False):
         for i, layer in enumerate(self.layers):
             layer_process = layer.process(processed_image)
             processed_image = layer_process[0]
-            topic_name = f"~/{self.instance_id if self.unique_name == '' else self.unique_name}/{layer.name}_{i}"
+            topic_name = f"/{self.instance_id if self.unique_name == '' else self.unique_name}/{layer.name}_{i}"
 
             if layer_process[1] is not None:
                 self.analysis_dict[f"{layer.name}_{i}"] = layer_process[1]
@@ -151,7 +154,7 @@ def run(self, in_image, verbose = False):
             if verbose: # Create a display showing how each layer processes the image before it
                 try:
                     # when debugging we expect different image encodings (maybe there's an RGB layer, then BW, etc.)
-                    verbose_layer_pub = Image_Publisher(topic_name, self)
+                    verbose_layer_pub = Image_Publisher(f"/front_cam/{layer.__class__.__name__}", self)
                     verbose_layer_pub.publish(processed_image)
                     ros_is_running = True
                 except:
@@ -162,15 +165,18 @@ def run(self, in_image, verbose = False):
 
                     if num_rows == 1:
                         axes[col_index].imshow(processed_image)
-                        axes[col_index].set_title(layer.name + "_" + i)
+                        axes[col_index].set_title(layer.name + "_" + str(i))
                     else:
                         axes[row_index, col_index].imshow(processed_image)
-                        axes[row_index, col_index].set_title(layer.name + "_" + i)                
+                        axes[row_index, col_index].set_title(layer.name + "_" + str(i))                
                     if num_rows == 1:
                         axes[col_index].axis('off')
                     else:
                         axes[row_index, col_index].axis('off')
 
+            else:
+                print("Not creating publisher")
+
         self.processed_image = processed_image
 
         if verbose and not ros_is_running:

vision_stack/layers/NormalizationLayer.py

@@ -33,7 +33,7 @@ def __init__(self) -> None:
         """
         Normalizes the image by calculating the median and the IQR of all the pixels, then subtracting the median form the pixels and dividing by the IQR.
         """
-        super().__init__()
+        super().__init__("RobustScalingLayer")
 
     def process(self, image):
         median = np.median(image)

vision_stack/layers/ResizeLayer.py

@@ -1,21 +1,21 @@
-import cv2
-from .Layer import PreprocessLayer
-
-class ResizeLayer(PreprocessLayer):
-
-    def __init__(self, new_width, new_height) -> None:
-        """
-        Resizes the image to be of new dimensions.
-
-        Parameters:
-            new_width: Image width will be resized to the new pixel width value.
-            new_height: Image height will be resized to the new pixel height value.
-        """
-        super().__init__("resize")
-        self.new_width = new_width
-        self.new_height = new_height
-
-
-    def process(self, image):
-        resized_image = cv2.resize(image, (self.new_width, self.new_height))
+import cv2
+from .Layer import PreprocessLayer
+
+class ResizeLayer(PreprocessLayer):
+
+    def __init__(self, new_width, new_height) -> None:
+        """
+        Resizes the image to be of new dimensions.
+
+        Parameters:
+            new_width: Image width will be resized to the new pixel width value.
+            new_height: Image height will be resized to the new pixel height value.
+        """
+        super().__init__("resize")
+        self.new_width = new_width
+        self.new_height = new_height
+
+    
+    def process(self, image):
+        resized_image = cv2.resize(image, (self.new_width, self.new_height))
         return (resized_image, None)

vision_stack/layers/UnderwaterEnhancementLayer.py

@@ -9,6 +9,7 @@
 
 import os
 import torch
+from pathlib import Path
 from PIL import Image
 import torch
 import numpy as np
@@ -18,7 +19,7 @@
 
 from .Layer import PreprocessLayer
 
-WEIGHTS_PATH = os.path.join(os.getcwd(), "vision_stack", "weights", 'model_best_2842.pth.tar')
+WEIGHTS_PATH = Path(os.path.abspath(__file__)).parent.parent / "weights" / "model_best_2842.pth.tar"
 
 class UnderWaterImageEnhancementLayer(PreprocessLayer):
     def __init__(self) -> None:
@@ -104,4 +105,4 @@ def forward(self,x):
         A = self.ANet(x)
         t = self.tNet(torch.cat((x*0+A,x),1))
         out = ((x-A)*t + A)
-        return torch.clamp(out,0.,1.)
+        return torch.clamp(out,0.,1.)