diff --git a/.gitignore b/.gitignore index ed514b8..eb9b2f2 100644 --- a/.gitignore +++ b/.gitignore @@ -13,5 +13,3 @@ __pycache__/ *.pyc *.pyo -# YOLO model (auto-downloaded on first run) -yolov8n.pt diff --git a/run_all.sh b/run_all.sh index d65843e..db93521 100755 --- a/run_all.sh +++ b/run_all.sh @@ -1,6 +1,11 @@ #!/usr/bin/env bash set -euo pipefail -export PYTHONPATH="/workspace/.site-packages:$PYTHONPATH" +# Link pip deps from EFS vendor if /workspace/.site-packages is missing +if [[ ! -d "${ROOT}/.site-packages" && -d "/opt/vendor/yolov8" ]]; then + ln -sf /opt/vendor/yolov8 "${ROOT}/.site-packages" +fi + +export PYTHONPATH="${ROOT}/.site-packages:$PYTHONPATH" export DISPLAY=:1 # Auto-download YOLO model if not present diff --git a/src/vision_yolo/vision_yolo/detector.py b/src/vision_yolo/vision_yolo/detector.py index 35a79ad..24eba8e 100644 --- a/src/vision_yolo/vision_yolo/detector.py +++ b/src/vision_yolo/vision_yolo/detector.py @@ -5,34 +5,42 @@ from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithP from cv_bridge import CvBridge from ultralytics import YOLO import cv2 +import time class YoloDetector(Node): def __init__(self): - super().__init__('yolo_detector') + super().__init__("yolo_detector") self.bridge = CvBridge() - self.get_logger().info('Loading YOLOv8n model...') - self.model = YOLO('yolov8n.pt') - self.get_logger().info('YOLOv8n model loaded') + self.get_logger().info("Loading YOLOv8n model...") + self.model = YOLO("yolov8n.pt") + self.get_logger().info("YOLOv8n model loaded") self.sub = self.create_subscription( - Image, '/camera/image_raw', self.on_image, 10) + Image, "/camera/image_raw", self.on_image, 10) self.pub = self.create_publisher( - Detection2DArray, '/vision/detections', 10) + Detection2DArray, "/vision/detections", 10) self.pub_vis = self.create_publisher( - Image, '/vision/image_annotated', 10) + Image, "/vision/image_annotated", 10) - self.get_logger().info('YOLO detector node started, subscribing to /camera/image_raw') + self.last_infer_time = 0.0 + self.infer_interval = 1.0 + self.get_logger().info("YOLO detector node started, subscribing to /camera/image_raw") def on_image(self, msg): + now = time.time() + if now - self.last_infer_time < self.infer_interval: + return + self.last_infer_time = now + try: - cv_img = self.bridge.imgmsg_to_cv2(msg, 'bgr8') - # CPU 降采样加速 - cv_img = cv2.resize(cv_img, (320, 240)) + cv_img = self.bridge.imgmsg_to_cv2(msg, "bgr8") + cv_img = cv2.resize(cv_img, (640, 480)) + cv2.imwrite("/workspace/camera_latest.jpg", cv_img) except Exception as e: - self.get_logger().error(f'cv_bridge failed: {e}') + self.get_logger().error(f"cv_bridge failed: {e}") return - results = self.model(cv_img, verbose=False, imgsz=320) + results = self.model(cv_img, verbose=False, imgsz=640) dets = Detection2DArray() dets.header = msg.header @@ -41,7 +49,6 @@ class YoloDetector(Node): cls_id = int(box.cls[0]) conf = float(box.conf[0]) label = self.model.names[cls_id] - x1, y1, x2, y2 = map(float, box.xyxy[0]) det = Detection2D() @@ -58,14 +65,15 @@ class YoloDetector(Node): dets.detections.append(det) cv2.rectangle(cv_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) - cv2.putText(cv_img, f'{label} {conf:.2f}', (int(x1), int(y1) - 10), + cv2.putText(cv_img, f"{label} {conf:.2f}", (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) self.pub.publish(dets) if dets.detections: - self.get_logger().info(f'Detected: {[d.results[0].hypothesis.class_id for d in dets.detections]}') + classes = [d.results[0].hypothesis.class_id for d in dets.detections] + self.get_logger().info(f"Detected: {classes}") - vis_msg = self.bridge.cv2_to_imgmsg(cv_img, 'bgr8') + vis_msg = self.bridge.cv2_to_imgmsg(cv_img, "bgr8") vis_msg.header = msg.header self.pub_vis.publish(vis_msg) @@ -79,5 +87,5 @@ def main(args=None): node.destroy_node() rclpy.shutdown() -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/yolov8n.pt b/yolov8n.pt new file mode 100644 index 0000000..0db4ca4 Binary files /dev/null and b/yolov8n.pt differ