Skip to content

Commit

Permalink
Merge branch 'master' of https://github.com/butia-bots/butia_vision i…
Browse files Browse the repository at this point in the history
…nto feature/vlm-recognition
  • Loading branch information
crislmfroes committed Aug 12, 2024
2 parents 06ac07b + 8a730b1 commit 99e4dc7
Show file tree
Hide file tree
Showing 9 changed files with 174 additions and 96 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
*.lua
*.csv
*.dat
*.pt
*.caffemodel
*.pt
*.pbtxt
*.prototxt
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,6 @@ def __mountLutTable(self, camera_info):
def __detectionDescriptionProcessing(self, data, description2d, header):
center_x, center_y = description2d.bbox.center.x, description2d.bbox.center.y
bbox_size_x, bbox_size_y = description2d.bbox.size_x, description2d.bbox.size_y
print(bbox_size_x, bbox_size_y)
if bbox_size_x == 0 or bbox_size_y == 0:
rospy.logwarn('BBox with zero size.')
return None
Expand Down Expand Up @@ -280,7 +279,14 @@ def __detectionDescriptionProcessing(self, data, description2d, header):
box_points = box_points[np.logical_and(np.all(box_points > min_bound, axis=1),
np.all(box_points < max_bound, axis=1))]

box = o3d.geometry.AxisAlignedBoundingBox().create_from_points(o3d.utility.Vector3dVector(box_points))
box_pcd = o3d.geometry.PointCloud()
box_pcd.points = o3d.utility.Vector3dVector(box_points)

box_pcd, _ = box_pcd.remove_statistical_outlier(20, 2.0)

# TODO: add a clustering algorithm and pick the closest cluster

box = o3d.geometry.AxisAlignedBoundingBox().create_from_points(box_pcd.points)

else:
box = o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound)
Expand Down Expand Up @@ -354,7 +360,6 @@ def __createDescription3D(self, source_data, description2d : Description2D, head
description3D : Description3D = self.DESCRIPTION_PROCESSING_ALGORITHMS[description2d.type](source_data, description2d, header)
description3D.bbox2D = description2d.bbox
description3D.class_num = description2d.class_num
print(description3D.class_num)
return description3D
else:
return None
Expand Down
162 changes: 81 additions & 81 deletions butia_recognition/config/yolov8_object_recognition.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -10,90 +10,90 @@ classes_by_category:
Animals: ['bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'teddybear']
Household: ['backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'vase', 'scissors', 'book', 'toothbrush']
all_classes:
- 'person'
- 'bicycle'
- 'car'
- 'motorcycle'
- 'airplane'
- 'bus'
- 'train'
- 'truck'
- 'boat'
- 'traffic light'
- 'fire hydrant'
- 'stop sign'
- 'parking meter'
- 'bench'
- 'bird'
- 'cat'
- 'dog'
- 'horse'
- 'sheep'
- 'cow'
- 'elephant'
- 'bear'
- 'zebra'
- 'giraffe'
- 'backpack'
- 'umbrella'
- 'handbag'
- 'tie'
- 'suitcase'
- 'frisbee'
- 'skis'
- 'snowboard'
- 'sports ball'
- 'kite'
- 'baseball bat'
- 'baseball glove'
- 'skateboard'
- 'surfboard'
- 'tennis racket'
- 'bottle'
- 'wine glass'
- 'cup'
- 'fork'
- 'knife'
- 'spoon'
- 'bowl'
- 'banana'
- 'apple'
- 'sandwich'
- 'orange'
- 'broccoli'
- 'carrot'
- 'hot dog'
- 'pizza'
- 'donut'
- 'cake'
- 'chair'
- 'couch'
- 'potted plant'
- 'bed'
- 'dining table'
- 'toilet'
- 'tv'
- 'laptop'
- 'mouse'
- 'remote'
- 'keyboard'
- 'cell phone'
- 'microwave'
- 'oven'
- 'toaster'
- 'sink'
- 'refrigerator'
- 'book'
- 'clock'
- 'vase'
- 'scissors'
- 'teddy bear'
- 'hair drier'
- 'toothbrush'
- 'person'
- 'bicycle'
- 'car'
- 'motorcycle'
- 'airplane'
- 'bus'
- 'train'
- 'truck'
- 'boat'
- 'traffic light'
- 'fire hydrant'
- 'stop sign'
- 'parking meter'
- 'bench'
- 'bird'
- 'cat'
- 'dog'
- 'horse'
- 'sheep'
- 'cow'
- 'elephant'
- 'bear'
- 'zebra'
- 'giraffe'
- 'backpack'
- 'umbrella'
- 'handbag'
- 'tie'
- 'suitcase'
- 'frisbee'
- 'skis'
- 'snowboard'
- 'sports ball'
- 'kite'
- 'baseball bat'
- 'baseball glove'
- 'skateboard'
- 'surfboard'
- 'tennis racket'
- 'bottle'
- 'wine glass'
- 'cup'
- 'fork'
- 'knife'
- 'spoon'
- 'bowl'
- 'banana'
- 'apple'
- 'sandwich'
- 'orange'
- 'broccoli'
- 'carrot'
- 'hot dog'
- 'pizza'
- 'donut'
- 'cake'
- 'chair'
- 'couch'
- 'potted plant'
- 'bed'
- 'dining table'
- 'toilet'
- 'tv'
- 'laptop'
- 'mouse'
- 'remote'
- 'keyboard'
- 'cell phone'
- 'microwave'
- 'oven'
- 'toaster'
- 'sink'
- 'refrigerator'
- 'book'
- 'clock'
- 'vase'
- 'scissors'
- 'teddy bear'
- 'hair drier'
- 'toothbrush'


max_sizes:
- [0.4, 2.5, 0.5]
- [0.05, 0.05, 0.05]

model_file: yolov8n.pt

Expand Down
59 changes: 59 additions & 0 deletions butia_recognition/config/yolov8_object_recognition_custom.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
threshold: 0.3
classes_by_category:
Drinks: ['nescau', 'kuat', 'cononut water', 'fanta']
Cleaning supplies: ['detergent', 'sponge', 'cloth']
Pantry items: ['gelatin', 'mustard', 'shoyo', 'sauce', 'tea']
Fruits: ['apple', 'pear', 'tangerine']
Snacks: ['treloso', 'chocolate', 'peanut']

all_classes:
- 'apple'
- 'chocolate'
- 'cloth'
- 'cononut water'
- 'detergent'
- 'fanta'
- 'gelatin'
- 'kuat'
- 'mustard'
- 'nescau'
- 'peanut'
- 'pear'
- 'sauce'
- 'shoyo'
- 'sponge'
- 'tangerine'
- 'tea'
- 'treloso'

max_sizes:
- [0.05, 0.05, 0.05]

model_file: yolov8-larc2023.pt

subscribers:

queue_size: 1
exact_time: false
slop: 0.2

image_rgb: /butia_vision/bvb/image_rgb
camera_info: /butia_vision/bvb/camera_info
image_depth: /butia_vision/bvb/image_depth

publishers:

object_recognition:
topic: /butia_vision/br/object_recognition
queue_size: 1

servers:

list_classes:
service: /butia_vision/br/object_recognition/list_classes

start:
service: /butia_vision/br/object_recognition/start

stop:
service: /butia_vision/br/object_recognition/stop
2 changes: 1 addition & 1 deletion butia_recognition/launch/yolov8_object_recognition.launch
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
<arg name="use_machine" default="true"/>
<arg name="output" default="screen"/>
<arg name="node_name" default="butia_object_recognition"/>
<arg name="config_file" default="yolov8_object_recognition.yaml"/>
<arg name="config_file" default="yolov8_object_recognition_custom.yaml"/>

<machine name="localhost" address="localhost" if="$(arg use_machine)"/>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@ def callback(self, *args):
img = ros_numpy.numpify(img)

debug_img = deepcopy(img)
debug_img = cv.cvtColor(debug_img, cv.COLOR_BGR2RGB)
results = None
bboxs = None

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from geometry_msgs.msg import Vector3
from butia_vision_msgs.msg import Description2D, Recognitions2D
import torch

import rospkg

class YoloV8Recognition(BaseRecognition):
def __init__(self, state=True):
Expand Down Expand Up @@ -41,7 +41,7 @@ def serverStop(self, req):
return super().serverStop(req)

def loadModel(self):
self.model = YOLO("yolov8n.pt")
self.model = YOLO(self.model_file)
self.model.conf = self.threshold
print('Done loading model!')

Expand Down Expand Up @@ -74,8 +74,7 @@ def callback(self, *args):
description_header = img_rgb.header
description_header.seq = 0

results = self.model.predict(cv_img)
debug_img = cv_img
results = self.model.predict(cv_img[:, :, ::-1], conf=self.threshold, verbose=True)
boxes_ = results[0].boxes.cpu().numpy()

if len(results[0].boxes):
Expand All @@ -87,14 +86,14 @@ def callback(self, *args):
continue

label_class = self.all_classes[int(box.cls)]

max_size = self.max_sizes[label_class]

description = Description2D()
description.header = copy(description_header)
description.type = Description2D.DETECTION
description.id = description.header.seq
description.score = float(box.conf)
description.max_size = Vector3(*[0.05, 0.05, 0.05])
description.max_size = Vector3(*max_size)
size = int(xyxy_box[2] - xyxy_box[0]), int(xyxy_box[3] - xyxy_box[1])
description.bbox.center.x = int(xyxy_box[0]) + int(size[0]/2)
description.bbox.center.y = int(xyxy_box[1]) + int(size[1]/2)
Expand All @@ -116,7 +115,7 @@ def callback(self, *args):
description.label = index + '/' + label_class if index is not None else label_class
objects_recognition.descriptions.append(description)

debug_img = results[0].plot()
debug_img = results[0].plot()[:, :, ::-1]
description_header.seq += 1

self.debug_publisher.publish(ros_numpy.msgify(Image, debug_img, 'rgb8'))
Expand All @@ -127,7 +126,7 @@ def callback(self, *args):
if len(people_recognition.descriptions) > 0:
self.people_detection_publisher.publish(people_recognition)
else:
debug_img = results[0].plot()
debug_img = results[0].plot()[:, :, ::-1]
self.debug_publisher.publish(ros_numpy.msgify(Image, debug_img, 'rgb8'))

def readParameters(self):
Expand All @@ -145,7 +144,21 @@ def readParameters(self):

self.all_classes = list(rospy.get_param("~all_classes", []))
self.classes_by_category = dict(rospy.get_param("~classes_by_category", {}))
self.model_file = rospy.get_param("~model_file", "yolov8_lab_objects.pt")

r = rospkg.RosPack()
r.list()
self.model_file = r.get_path("butia_recognition") + "/weigths/" + rospy.get_param("~model_file", "yolov8n.pt")

max_sizes = rospy.get_param("~max_sizes", [[0.05, 0.05, 0.05]])

if len(max_sizes) == 1:
max_sizes = [max_sizes[0] for _ in self.all_classes]

if len(max_sizes) != len(self.all_classes):
rospy.logerr('Wrong definition of max_sizes in yaml of recognition node.')
assert False

self.max_sizes = dict([(self.all_classes[i], max_sizes[i]) for i in range(len(max_sizes))])

super().readParameters()

Expand Down
Binary file removed butia_recognition/weigths/tmpoxfg42bv
Binary file not shown.
2 changes: 1 addition & 1 deletion butia_vision_msgs

0 comments on commit 99e4dc7

Please sign in to comment.