Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions test/topics_test.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -62,4 +65,4 @@ def custom_process(img, args):
rclpy.shutdown()



test_topic_creation()
73 changes: 73 additions & 0 deletions test/vision_stack_test.py
Original file line number Diff line number Diff line change
@@ -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()
13 changes: 7 additions & 6 deletions vision_stack/VisionStack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand All @@ -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:
Expand All @@ -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:
Expand Down
40 changes: 20 additions & 20 deletions vision_stack/layers/ResizeLayer.py
Original file line number Diff line number Diff line change
@@ -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)