Skip to content

Instantly share code, notes, and snippets.

@thomasweng15
Created March 21, 2022 15:52
Show Gist options
  • Save thomasweng15/63f371952bdac8cfaadd1c695012ac69 to your computer and use it in GitHub Desktop.
Save thomasweng15/63f371952bdac8cfaadd1c695012ac69 to your computer and use it in GitHub Desktop.
Segment HSV from an image rostopic using OpenCV sliders
import cv2
import rospy
import argparse
import numpy as np
import matplotlib.pyplot as plt
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
class Segmenter():
def __init__(self, args):
self.args = args
rospy.init_node("fg_bg_segmenter")
self.sub = rospy.Subscriber(args.img_topic, Image, self.cb)
self.title_window = 'Foreground-Background HSV Segmenter'
self.lower_hue = 0
self.upper_hue = 255
self.lower_sat = 0
self.upper_sat = 255
self.lower_val = 0
self.upper_val = 255
self.bgr = None
self.hsv = None
self.dst = None
self.bridge = CvBridge()
cv2.namedWindow(self.title_window)
cv2.createTrackbar('lower hue', self.title_window, 0, 255, self.on_lower_h)
cv2.createTrackbar('upper hue', self.title_window, 255, 255, self.on_upper_h)
cv2.createTrackbar('lower sat', self.title_window, 0, 255, self.on_lower_s)
cv2.createTrackbar('upper sat', self.title_window, 255, 255, self.on_upper_s)
cv2.createTrackbar('lower val', self.title_window, 0, 255, self.on_lower_v)
cv2.createTrackbar('upper val', self.title_window, 255, 255, self.on_upper_v)
def update(self):
lower = np.array([self.lower_hue, self.lower_sat, self.lower_val])
upper = np.array([self.upper_hue, self.upper_sat, self.upper_val])
mask = cv2.inRange(self.hsv, lower, upper)
# kernel = np.ones((9,9),np.uint8)
# mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
self.dst = cv2.bitwise_and(np.stack([mask, mask, mask], axis=2), self.bgr)
print("lower HSV: {}, upper HSV: {}".format(lower, upper))
def on_lower_h(self, val):
self.lower_hue = val
self.update()
def on_upper_h(self, val):
self.upper_hue = val
self.update()
def on_lower_s(self, val):
self.lower_sat = val
self.update()
def on_upper_s(self, val):
self.upper_sat = val
self.update()
def on_lower_v(self, val):
self.lower_val = val
self.update()
def on_upper_v(self, val):
self.upper_val = val
self.update()
def cb(self, msg):
im = self.bridge.imgmsg_to_cv2(msg)
self.bgr = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
self.hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
def spin(self):
while not rospy.is_shutdown():
if self.bgr is not None:
self.update()
if self.dst is None:
rospy.sleep(0.1)
else:
cv2.imshow(self.title_window, self.dst)
cv2.waitKey(30)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--img_topic', help='ROS topic to subscribe to for RGB image', default='/camera/color/image_rect_color')
args, _ = parser.parse_known_args()
# path = "/home/tweng/ws1/src/bimanual_folding/bimanual_ros/scripts/goals/real_rect/ms_two_side_horz_0_gray.png"
# path = "/home/tweng/ws1/src/bimanual_folding/bimanual_ros/scripts/goals/real_lblue/towel_train_9_high_gray.png"
# path = "/home/tweng/ws1/src/bimanual_folding/bimanual_ros/scripts/goals/real_shirt/open_2side_high_gray.png"
path = "/home/tweng/ws1/src/bimanual_folding/bimanual_ros/scripts/goals/real_high/open_2side_high_gray.png"
s = Segmenter(args)
s.spin()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment