diff --git a/test/topics_test.py b/test/topics_test.py index 623654d..c5cbf5c 100644 --- a/test/topics_test.py +++ b/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() diff --git a/test/vision_stack_test.py b/test/vision_stack_test.py new file mode 100644 index 0000000..a9fcbd3 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/vision_stack/VisionStack.py b/vision_stack/VisionStack.py index e229e7b..62809a4 100644 --- a/vision_stack/VisionStack.py +++ b/vision_stack/VisionStack.py @@ -10,9 +10,9 @@ 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 +32,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 @@ -135,7 +136,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 +152,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"{layer.__class__.__name__}", self) verbose_layer_pub.publish(processed_image) ros_is_running = True except: @@ -162,15 +163,15 @@ 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') - + self.processed_image = processed_image if verbose and not ros_is_running: diff --git a/vision_stack/layers/ResizeLayer.py b/vision_stack/layers/ResizeLayer.py index 50fb763..16f1a4a 100644 --- a/vision_stack/layers/ResizeLayer.py +++ b/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) \ No newline at end of file