From ff08f4d71cb0491b5e31fbe6119344076501b56c Mon Sep 17 00:00:00 2001 From: Peter van Dooren Date: Thu, 19 Oct 2023 12:31:55 +0200 Subject: [PATCH] move model into class --- ed_sensor_integration/scripts/listener.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ed_sensor_integration/scripts/listener.py b/ed_sensor_integration/scripts/listener.py index f420c3d4..202ebc2f 100755 --- a/ed_sensor_integration/scripts/listener.py +++ b/ed_sensor_integration/scripts/listener.py @@ -45,6 +45,11 @@ class table_segmentor: def __init__(self) -> None: + model_path = "~/MEGA/developers/Donal/yolov8x-seg.pt" + device = "cuda" + self.model = YOLO(model_path).to(device) + self.table_class = 60 #table class defined with index 60 (person = 0) + rospy.init_node('listener', anonymous=True) self.publisher = rospy.Publisher('/hero/segmented_image',Image) self.subscriber = rospy.Subscriber('/hero/head_rgbd_sensor/rgb/image_raw',Image , self.callback) @@ -63,16 +68,11 @@ def callback(self, data): bridge = CvBridge() cv_image = bridge.imgmsg_to_cv2(data, desired_encoding='passthrough') rospy.loginfo("converted message") - #Yolov8 model code - model_path = "~/MEGA/developers/Donal/yolov8x-seg.pt" - device = "cuda" - model = YOLO(model_path).to(device) - table_class = 60 #table class defined with index 60 (person = 0) - classes, segmentations = self.detect(model, cv_image) + classes, segmentations = self.detect(self.model, cv_image) #extract table segment and add to frame for class_id, seg in zip(classes, segmentations): - if class_id == table_class: + if class_id == self.table_class: cv2.polylines(cv_image, [seg], True, (255,0,0), 2) # cv2.imshow("Segmented Image", cv_image) # cv2.waitKey(1)