1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
|
"""
CARLA RGB 图像 -> rf-DETR TensorRT 检测 -> OC-SORT 跟踪 + 移动轨迹 -> 发布画框图像 + Detection2DArray。
属于 perception_camera 包。用系统 python3 运行（rclpy/cv_bridge/tensorrt/torch/filterpy 在系统 python，不碰 conda）。
source install/setup.bash
ros2 run perception_camera carla_detector.py
订阅: /carla/hero/rgb/image (sensor_msgs/Image)
发布: /carla/hero/rgb/detection_image (画框+ID+轨迹, 给 rviz2)
/carla/hero/detections (vision_msgs/Detection2DArray, track_id 在 Detection2D.id)
>>> 使用前替换下面两个路径为你自己的 <<<
注意（踩过的坑，已修复）：
- COCO 类名: 模型 argmax 直接 = COCO ID，不要 +1（+1 会把 car 标成 motorcycle）
- 类别过滤: ALLOWED_CLASSES 必须用 list，np.isin 对 set 返回全 False
- TRT 上下文变量名不能叫 self.context（撞 rclpy Node 只读属性），用 self.trt_context
- humble 的 BoundingBox2D.center 是 vision_msgs/Pose2D，用 center.position.x（不是 center.x）
"""
import os
import sys
import numpy as np
import cv2
import torch
import tensorrt as trt
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSHistoryPolicy
from sensor_msgs.msg import Image
from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose
from cv_bridge import CvBridge
from collections import OrderedDict, namedtuple, defaultdict
ENGINE_PATH = os.path.expanduser("~/models/rf-detr/weights/onnx/rfdetr-medium.engine")
OC_SORT_PATH = os.path.expanduser("~/models/OC_SORT")
sys.path.insert(0, OC_SORT_PATH)
from trackers.ocsort_tracker.ocsort import OCSort
INPUT_TOPIC = "/carla/hero/rgb/image"
OUT_IMAGE_TOPIC = "/carla/hero/rgb/detection_image"
OUT_DET_TOPIC = "/carla/hero/detections"
RESOLUTION = 576
THRESHOLD = 0.5
ALLOWED_CLASSES = [1, 2, 3, 4, 6, 8]
TRAIL_LEN = 40
MEAN = torch.tensor([0.485, 0.456, 0.406])
STD = torch.tensor([0.229, 0.224, 0.225])
COCO_CLASSES = {
1: "person", 2: "bicycle", 3: "car", 4: "motorcycle", 5: "airplane", 6: "bus",
7: "train", 8: "truck", 9: "boat", 10: "traffic light", 11: "fire hydrant",
13: "stop sign", 14: "parking meter", 15: "bench", 16: "bird", 17: "cat",
18: "dog", 19: "horse", 20: "sheep", 21: "cow", 22: "elephant", 23: "bear",
24: "zebra", 25: "giraffe", 27: "backpack", 28: "umbrella", 31: "handbag",
32: "tie", 33: "suitcase", 34: "frisbee", 35: "skis", 36: "snowboard",
37: "sports ball", 38: "kite", 39: "baseball bat", 40: "baseball glove",
41: "skateboard", 42: "surfboard", 43: "tennis racket", 44: "bottle",
46: "wine glass", 47: "cup", 48: "fork", 49: "knife", 50: "spoon", 51: "bowl",
52: "banana", 53: "apple", 54: "sandwich", 55: "orange", 56: "broccoli",
57: "carrot", 58: "hot dog", 59: "pizza", 60: "donut", 61: "cake", 62: "chair",
63: "couch", 64: "potted plant", 65: "bed", 67: "dining table", 70: "toilet",
72: "tv", 73: "laptop", 74: "mouse", 75: "remote", 76: "keyboard",
77: "cell phone", 78: "microwave", 79: "oven", 80: "toaster", 81: "sink",
82: "refrigerator", 84: "book", 85: "clock", 86: "vase", 87: "scissors",
88: "teddy bear", 89: "hair drier", 90: "toothbrush",
}
_TRT_DTYPE_TO_NP = {
trt.DataType.FLOAT: np.float32, trt.DataType.HALF: np.float16,
trt.DataType.INT32: np.int32, trt.DataType.INT64: np.int64,
trt.DataType.INT8: np.int8, trt.DataType.UINT8: np.uint8,
trt.DataType.BOOL: np.bool_,
}
def _iou_one_to_many(box, boxes):
"""box: [x1,y1,x2,y2]; boxes: (N,4)。返回每个的 IoU。"""
if len(boxes) == 0:
return np.zeros(0)
x1 = np.maximum(box[0], boxes[:, 0])
y1 = np.maximum(box[1], boxes[:, 1])
x2 = np.minimum(box[2], boxes[:, 2])
y2 = np.minimum(box[3], boxes[:, 3])
inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
area1 = (box[2] - box[0]) * (box[3] - box[1])
area2 = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
return inter / (area1 + area2 - inter + 1e-9)
def _color_for(tid):
"""根据 track_id 生成稳定颜色 (RGB)。"""
return (int((tid * 53) % 256), int((tid * 97) % 256), int((tid * 151) % 256))
class CarlaDetector(Node):
def __init__(self):
super().__init__("carla_detector")
self.bridge = CvBridge()
self._frame_count = 0
self._load_engine()
self.tracker = OCSort(det_thresh=THRESHOLD)
self.trails = defaultdict(list)
qos = QoSProfile(
reliability=QoSReliabilityPolicy.BEST_EFFORT,
history=QoSHistoryPolicy.KEEP_LAST, depth=5,
)
self.sub = self.create_subscription(Image, INPUT_TOPIC, self.image_cb, qos)
self.pub_img = self.create_publisher(Image, OUT_IMAGE_TOPIC, 10)
self.pub_det = self.create_publisher(Detection2DArray, OUT_DET_TOPIC, 10)
self.get_logger().info(
f"carla_detector ready (detection+OC-SORT). "
f"sub={INPUT_TOPIC} pub_img={OUT_IMAGE_TOPIC} pub_det={OUT_DET_TOPIC}"
)
def _load_engine(self):
Binding = namedtuple("Binding", ("name", "dtype", "shape", "data"))
logger = trt.Logger(trt.Logger.WARNING)
with open(ENGINE_PATH, "rb") as f:
self.engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())
self.trt_context = self.engine.create_execution_context()
self.bindings = OrderedDict()
for name in self.engine:
shape = self.engine.get_tensor_shape(name)
dtype = _TRT_DTYPE_TO_NP[self.engine.get_tensor_dtype(name)]
data = torch.from_numpy(np.empty(shape, dtype=dtype)).cuda()
self.bindings[name] = Binding(name, dtype, shape, data)
self.input_names = [n for n in self.bindings
if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT]
self.binding_addrs = [int(self.bindings[n].data.data_ptr()) for n in self.bindings]
self.get_logger().info(f"engine loaded: inputs={self.input_names}")
def _preprocess(self, img_rgb):
h, w = img_rgb.shape[:2]
t = torch.from_numpy(img_rgb).float().permute(2, 0, 1) / 255.0
t = torch.nn.functional.interpolate(
t.unsqueeze(0), size=(RESOLUTION, RESOLUTION),
mode="bilinear", align_corners=False).squeeze(0)
t[0] = (t[0] - MEAN[0]) / STD[0]
t[1] = (t[1] - MEAN[1]) / STD[1]
t[2] = (t[2] - MEAN[2]) / STD[2]
return t.unsqueeze(0).cuda(), (w, h)
def _infer(self, inp):
self.bindings[self.input_names[0]].data.copy_(inp)
self.trt_context.execute_v2(self.binding_addrs)
torch.cuda.synchronize()
dets = self.bindings["dets"].data.cpu().numpy()
labels = self.bindings["labels"].data.cpu().numpy()
return dets, labels
def _postprocess(self, dets, labels, orig_w, orig_h):
logits = labels[0, :, :-1]
scores_all = 1.0 / (1.0 + np.exp(-logits.clip(-88, 88)))
scores = scores_all.max(axis=-1)
class_ids = scores_all.argmax(axis=-1)
keep = scores > THRESHOLD
keep = np.logical_and(keep, np.isin(class_ids, ALLOWED_CLASSES))
cxcywh = dets[0][keep]
scores = scores[keep]
class_ids = class_ids[keep]
cx, cy, bw, bh = cxcywh[:, 0], cxcywh[:, 1], cxcywh[:, 2], cxcywh[:, 3]
x1 = (cx - bw / 2.0) * orig_w
y1 = (cy - bh / 2.0) * orig_h
x2 = (cx + bw / 2.0) * orig_w
y2 = (cy + bh / 2.0) * orig_h
boxes = np.stack([x1, y1, x2, y2], axis=-1)
return boxes, scores, class_ids
def image_cb(self, msg):
try:
img = self.bridge.imgmsg_to_cv2(msg, desired_encoding="rgb8")
except Exception as e:
self.get_logger().warn(f"cv_bridge convert failed: {e}")
return
inp, (w, h) = self._preprocess(img)
dets_raw, labels = self._infer(inp)
boxes, scores, cids = self._postprocess(dets_raw, labels, w, h)
if len(boxes) > 0:
dets_in = np.concatenate([boxes, scores.reshape(-1, 1)], axis=1)
else:
dets_in = np.empty((0, 5))
tracks = self.tracker.update(dets_in, (h, w), (h, w))
ann = img.copy()
da = Detection2DArray()
da.header = msg.header
for trk in tracks:
x1, y1, x2, y2, tid = trk
x1i, y1i, x2i, y2i = int(x1), int(y1), int(x2), int(y2)
tid = int(tid)
cls_name = "obj"
best, best_iou = -1, 0.0
if len(boxes) > 0:
ious = _iou_one_to_many(np.array([x1, y1, x2, y2]), boxes)
best = int(ious.argmax())
best_iou = float(ious[best])
if best_iou > 0.3:
coco_id = int(cids[best])
cls_name = COCO_CLASSES.get(coco_id, f"id{coco_id}")
color = _color_for(tid)
cx, cy = (x1i + x2i) // 2, (y1i + y2i) // 2
self.trails[tid].append((cx, cy))
if len(self.trails[tid]) > TRAIL_LEN:
self.trails[tid].pop(0)
pts = self.trails[tid]
for i in range(1, len(pts)):
cv2.line(ann, pts[i - 1], pts[i], color, 2)
cv2.rectangle(ann, (x1i, y1i), (x2i, y2i), color, 2)
cv2.putText(ann, f"{cls_name}:{tid}", (x1i, max(0, y1i - 6)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
d = Detection2D()
d.header = msg.header
d.id = str(tid)
d.bbox.center.position.x = float((x1 + x2) / 2.0)
d.bbox.center.position.y = float((y1 + y2) / 2.0)
d.bbox.size_x = float(x2 - x1)
d.bbox.size_y = float(y2 - y1)
if best >= 0 and best_iou > 0.3:
hyp = ObjectHypothesisWithPose()
hyp.hypothesis.class_id = str(int(cids[best]))
hyp.hypothesis.score = float(scores[best])
d.results.append(hyp)
da.detections.append(d)
out_msg = self.bridge.cv2_to_imgmsg(ann, encoding="rgb8")
out_msg.header = msg.header
self.pub_img.publish(out_msg)
self.pub_det.publish(da)
self._frame_count += 1
if self._frame_count % 30 == 0:
self.get_logger().info(
f"frame {self._frame_count}: {len(boxes)} dets, {len(tracks)} tracks")
def main():
rclpy.init()
node = CarlaDetector()
try:
rclpy.spin(node)
except KeyboardInterrupt:
pass
node.destroy_node()
rclpy.shutdown()
if __name__ == "__main__":
main()
|
