Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add python library for yolov7 node #2

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
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
4 changes: 4 additions & 0 deletions src/detect_ros.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from vision_msgs.msg import Detection2DArray, Detection2D, BoundingBox2D
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from detection_server import DetectionServer
import yaml


Expand Down Expand Up @@ -152,6 +153,9 @@ def process_img_msg(self, img_msg: Image):
# publishing
detection_msg = create_detection_msg(img_msg, detections)
self.detection_publisher.publish(detection_msg)
# post to detection library
DetectionServer.update_detections(np_img_resized, detections)


# visualizing if required
if self.visualization_publisher:
Expand Down
82 changes: 82 additions & 0 deletions src/detection_server.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
import rospy
import torch
from numpy import ndarray
from utils.Detection2DArray import Detection2DArray
from utils.Detection2D import Detection2D
from utils.BoundingBox2D import BoundingBox2D
from utils.ObjectHypothesisWithPose import ObjectHypothesisWithPose

# todo: update to classes
class DetectionServer:
def __init__(self):
self.detection_array = Detection2DArray()
self.detection_history = []
self.number_of_detections = 0
self.objects_found = False

# todo: find our if I still need this
def create_time_stamp():
h = Header()
h.stamp = rospy.Time.now()
return h


def update_detections(image: ndarray, detections: torch.Tensor) -> Detection2DArray:
"""
:param detections: torch tensor of shape [num_boxes, 6] where each element is
[x1, y1, x2, y2, confidence, class_id]
:returns: detections as a ros message of type Detection2DArray
"""
self.detection_array = Detection2DArray()

# todo: add source image https://docs.ros.org/en/lunar/api/vision_msgs/html/msg/Detection2D.html

# header
time_stamp = create_time_stamp()
self.detection_array.header = time_stamp
self.detection_array.img = image
for detection in detections:
x1, y1, x2, y2, conf, cls = detection.tolist()
single_detection_msg = Detection2D()
single_detection_msg.header = time_stamp


# bbox
bbox = BoundingBox2D()
w = int(round(x2 - x1))
h = int(round(y2 - y1))
cx = int(round(x1 + w / 2))
cy = int(round(y1 + h / 2))
bbox.size_x = w
bbox.size_y = h
bbox.center_x = cx
bbox.center_y = cy

single_detection_msg.bbox = bbox

# class id & confidence
obj_hyp = ObjectHypothesisWithPose()
obj_hyp.id = int(cls)
obj_hyp.score = conf
single_detection_msg.results = [obj_hyp]

self.detection_array.detections.append(single_detection_msg)
self.detection_history.append(self.detection_array)
return self.detection_array

def get_detections():
if(len(self.detection_array)==0:
raise Exception("Detections were requested but no objects were found")
else:
return self.detection_array

def get_number_of_objects_detected():
return len(self.detection_array)

def get_detection_history():
return self.detection_history

def erase_detection_history():
self.detection_history = []


7 changes: 7 additions & 0 deletions src/utils/BoundingBox2D.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
class BoundingBox2D:
def __init__(self, cx: int, cy: int, sx: int, sy: int):
self.center_x = cx
self.center_y = cy
self.size_x = sx
self.size_y = sy

8 changes: 8 additions & 0 deletions src/utils/Detection2D.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
from ObjectHypothesisWithPose import ObjectHypothesisWithPose
from BoundingBox1D import BoundingBox2D

class Detection2D:
def __init__(self, results: ObjectHypothesisWithPose, bbox: BoundingBox2D, header):
self.results = results
self.bbox = bbox
self.header = header
7 changes: 7 additions & 0 deletions src/utils/Detection2DArray.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
from Detection2D import Detection2D
from numpy import ndarray
class Detection2DArray:
def __init__(self, img: ndarray, detections: list[Detection2D], header):
self.img = img
self.detections = detections
self.header = header
5 changes: 5 additions & 0 deletions src/utils/ObjectHypothesisWithPose.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
class ObjectHypothesisWithPose:
def __init__(self, id: int, score: float):
self.id = id
self.score = score