VLM added
This commit is contained in:
Bug
parent
c209618f83
commit
c582d87725
|
|
@ -0,0 +1,173 @@
|
|||
# VLM GPU 视觉识别集成方案
|
||||
|
||||
> 在保留现有 YOLOv8 CPU 识别的基础上，增加运行时选项启用 GPU 本地 VLM 模型，用于 Gazebo office 场景机器人巡检与摄像头目标识别。
|
||||
|
||||
---
|
||||
|
||||
## 一、模型选型
|
||||
|
||||
| 模型 | 参数量 | 显存需求 | 推理速度* | 中文 | 视觉定位 | 许可证 |
|
||||
|------|--------|---------|----------|------|---------|--------|
|
||||
| **Qwen2.5-VL-3B-Instruct** ⭐ | 3B | ~8-10GB (BF16) | 2-4s/帧 | 原生 | ✅ bbox | Apache-2.0 |
|
||||
| Qwen2.5-VL-3B-AWQ | 3B | ~3-4GB (INT4) | 1.5-3s/帧 | 原生 | ✅ bbox | Apache-2.0 |
|
||||
| MiniCPM-V 2.6 | 8B | ~17GB / ~7GB(INT4) | 3-6s/帧 | 原生 | ✅ bbox | 学术/商用 |
|
||||
|
||||
*推理速度：单张 640×480 图像，生成 80-120 tokens，NVIDIA T4 估算值
|
||||
|
||||
**推荐：Qwen2.5-VL-3B-Instruct (BF16)**
|
||||
- 3B 参数性价比最高，T4 16GB 刚好跑得动
|
||||
- 原生支持中文视觉定位（`<|box_start|>(x1,y1),(x2,y2)<|box_end|>`）
|
||||
- 支持结构化 JSON 输出，便于解析为 ROS2 Detection2DArray
|
||||
- Apache-2.0 许可证，无商用限制
|
||||
|
||||
---
|
||||
|
||||
## 二、三种运行时模式
|
||||
|
||||
```bash
|
||||
# 模式 A：YOLOv8 CPU（默认，完全兼容原有行为）
|
||||
./run_all.sh mapping --explore --vision
|
||||
|
||||
# 模式 B：VLM GPU（异步推理，避免阻塞 Nav2）
|
||||
VISION_BACKEND=vlm ./run_all.sh mapping --explore --vision
|
||||
|
||||
# 模式 C：混合模式 — YOLO 持续检测 + VLM 每 10s 深度分析
|
||||
VISION_BACKEND=hybrid VLM_INTERVAL=10.0 ./run_all.sh mapping --explore --vision
|
||||
```
|
||||
|
||||
### 环境变量
|
||||
|
||||
| 变量 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `VISION_BACKEND` | `yolo` | `yolo` / `vlm` / `hybrid` |
|
||||
| `VLM_INTERVAL` | `5.0` | VLM 采样间隔（秒） |
|
||||
| `VLM_MODEL` | `Qwen/Qwen2.5-VL-3B-Instruct` | 模型名称 |
|
||||
| `VLM_MAX_NEW_TOKENS` | `256` | 最大生成 token 数 |
|
||||
| `VLM_PROMPT` | `（内置）` | 自定义提示词 |
|
||||
| `VLM_MIN_PIXELS` | `200704` | 图像最小像素（256×28×28） |
|
||||
| `VLM_MAX_PIXELS` | `501760` | 图像最大像素（640×28×28） |
|
||||
|
||||
---
|
||||
|
||||
## 三、安装部署
|
||||
|
||||
### 3.1 安装依赖（GPU 容器内执行）
|
||||
|
||||
```bash
|
||||
bash scripts/install_vlm.sh
|
||||
```
|
||||
|
||||
该脚本会自动：
|
||||
- 检测 CUDA 环境，按需安装 PyTorch 2.1.2+cu121
|
||||
- 安装 transformers、accelerate、qwen-vl-utils
|
||||
- Ampere 架构 GPU 自动安装 Flash Attention 2
|
||||
- 预下载 Qwen2.5-VL-3B-Instruct 模型（约 6-8GB，缓存到 EFS）
|
||||
|
||||
### 3.2 验证安装
|
||||
|
||||
```bash
|
||||
python3 -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"
|
||||
```
|
||||
|
||||
### 3.3 启动测试
|
||||
|
||||
```bash
|
||||
# YOLO 模式（默认）
|
||||
./run_all.sh mapping --explore --vision
|
||||
|
||||
# VLM 模式
|
||||
VISION_BACKEND=vlm VLM_INTERVAL=5.0 ./run_all.sh mapping --explore --vision
|
||||
|
||||
# 查看 VLM 场景描述
|
||||
ros2 topic echo /vision/scene_description
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 四、硬件与成本
|
||||
|
||||
### 推荐 GPU 实例（AWS 中国）
|
||||
|
||||
| 实例 | GPU | 显存 | 按需/小时 | 适用性 |
|
||||
|------|-----|------|----------|--------|
|
||||
| g4dn.xlarge | T4 | 16GB | ~¥3.8 | ✅ 最佳性价比 |
|
||||
| g5.xlarge | A10G | 24GB | ~¥7.2 | ✅ 性能冗余 |
|
||||
|
||||
### 显存占用
|
||||
|
||||
| 组件 | 显存 |
|
||||
|------|------|
|
||||
| Qwen2.5-VL-3B BF16 权重 | ~6.5GB |
|
||||
| ViT 视觉编码器 | ~1.5GB |
|
||||
| KV Cache | ~0.5GB |
|
||||
| 激活/中间结果 | ~1-2GB |
|
||||
| **总计** | **~10-11GB** |
|
||||
|
||||
> 若显存不足，可切换至 AWQ 量化版：`export VLM_MODEL=Qwen/Qwen2.5-VL-3B-Instruct-AWQ`
|
||||
|
||||
---
|
||||
|
||||
## 五、架构要点
|
||||
|
||||
### VLM 异步推理
|
||||
|
||||
VLM 单帧推理需 2-5 秒，若在 ROS2 回调中同步执行会阻塞 Nav2。`detector.py` 在 VLM 模式下自动使用**后台线程**执行推理，ROS executor 不会被阻塞。
|
||||
|
||||
### 双格式解析
|
||||
|
||||
VLM 输出支持两种解析策略：
|
||||
1. **JSON 提取**：优先匹配 `{"objects": [...], "scene_description": "..."}`
|
||||
2. **原生 grounding 兜底**：解析 `<|box_start|>(x1,y1),(x2,y2)<|box_end|>`
|
||||
|
||||
### 模型缓存
|
||||
|
||||
- 首次下载约 6-8GB，缓存到 `/workspace/.cache/huggingface`（EFS 持久化）
|
||||
- 容器重启后无需重新下载
|
||||
- 可在镜像构建时预置到 `/opt/vendor/vlm_models/` 以加速启动
|
||||
|
||||
---
|
||||
|
||||
## 六、性能预期
|
||||
|
||||
| 指标 | YOLOv8 CPU | VLM GPU (T4) |
|
||||
|------|-----------|-------------|
|
||||
| 单帧处理 | ~80ms | ~3-5s |
|
||||
| 频率 | 1Hz | 0.2Hz (5s interval) |
|
||||
| CPU | ~100% (1核) | ~50% (预处理) |
|
||||
| GPU | 0 | ~90% (推理时) |
|
||||
| 假阳性 | 高（COCO 预训练） | 低（语义理解） |
|
||||
| 语义能力 | 无 | 场景描述、异常检测 |
|
||||
|
||||
---
|
||||
|
||||
## 七、文件清单
|
||||
|
||||
| 文件 | 说明 |
|
||||
|------|------|
|
||||
| `src/vision_yolo/vision_yolo/detector.py` | 修改后的检测器（支持三模式切换） |
|
||||
| `scripts/install_vlm.sh` | VLM 依赖安装脚本 |
|
||||
| `docs/VLM_INTEGRATION.md` | 本文档 |
|
||||
|
||||
---
|
||||
|
||||
## 八、快速验证流程
|
||||
|
||||
```bash
|
||||
# 1. SSH 进入 GPU 实例
|
||||
ssh devuser@<gpu-space>
|
||||
|
||||
# 2. 安装依赖
|
||||
bash scripts/install_vlm.sh
|
||||
|
||||
# 3. 启动 Gazebo + Nav2
|
||||
./run_all.sh mapping --explore
|
||||
|
||||
# 4. 单独测试 VLM 节点
|
||||
export VISION_BACKEND=vlm
|
||||
export VLM_INTERVAL=3.0
|
||||
ros2 run vision_yolo detector
|
||||
|
||||
# 5. 观察话题
|
||||
ros2 topic echo /vision/scene_description
|
||||
ros2 topic hz /vision/detections
|
||||
```
|
||||
|
|
@ -306,6 +306,15 @@ fi
|
|||
|
||||
echo
|
||||
if [[ "$WITH_VISION" == true ]]; then
|
||||
# VLM 运行时配置（VISION_BACKEND=yolo|vlm|hybrid）
|
||||
export VISION_BACKEND="${VISION_BACKEND:-yolo}"
|
||||
export VLM_INTERVAL="${VLM_INTERVAL:-5.0}"
|
||||
export VLM_MODEL="${VLM_MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
export VLM_MAX_NEW_TOKENS="${VLM_MAX_NEW_TOKENS:-256}"
|
||||
export VLM_PROMPT="${VLM_PROMPT:-}"
|
||||
export VLM_MIN_PIXELS="${VLM_MIN_PIXELS:-200704}"
|
||||
export VLM_MAX_PIXELS="${VLM_MAX_PIXELS:-501760}"
|
||||
export HF_HOME="${HF_HOME:-/workspace/.cache/huggingface}"
|
||||
start_bg "D_vision_yolo" ros2 run vision_yolo detector
|
||||
last_index=$(( ${#PIDS[@]} - 1 ))
|
||||
VISION_PID="${PIDS[$last_index]}"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,166 @@
|
|||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
# =============================================================================
|
||||
# VLM (Qwen2.5-VL-3B) 依赖安装脚本
|
||||
# =============================================================================
|
||||
# 运行环境：ECS GPU 容器（/workspace 挂载在 EFS 上持久化）
|
||||
# 前置条件：NVIDIA GPU 可用，CUDA 驱动已安装
|
||||
# =============================================================================
|
||||
|
||||
ROOT="/workspace"
|
||||
SITE_PACKAGES="${ROOT}/.site-packages"
|
||||
VLM_MODEL_CACHE="${ROOT}/.cache/huggingface"
|
||||
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
NC='\033[0m'
|
||||
|
||||
log_info() { echo -e "${GREEN}[INFO]${NC} $*"; }
|
||||
log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
|
||||
log_error() { echo -e "${RED}[ERROR]${NC} $*"; }
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 0. 环境检查
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "检查 GPU / CUDA 环境..."
|
||||
|
||||
if ! command -v nvidia-smi &>/dev/null; then
|
||||
log_error "nvidia-smi 未找到。请确认当前为 GPU 实例且 NVIDIA 驱动已安装。"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader || true
|
||||
|
||||
if ! python3 -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
|
||||
log_warn "PyTorch CUDA 未安装或不可用，将重新安装 PyTorch CUDA 版..."
|
||||
NEED_PYTORCH_CUDA=1
|
||||
else
|
||||
python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}, Device: {torch.cuda.get_device_name(0)}')"
|
||||
NEED_PYTORCH_CUDA=0
|
||||
fi
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 1. 安装 PyTorch CUDA (如需要)
|
||||
# -----------------------------------------------------------------------------
|
||||
if [[ "$NEED_PYTORCH_CUDA" == "1" ]]; then
|
||||
log_info "安装 PyTorch 2.1.2 + CUDA 12.1 ..."
|
||||
pip install --target="$SITE_PACKAGES" --upgrade \
|
||||
torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 \
|
||||
--index-url https://download.pytorch.org/whl/cu121 \
|
||||
--no-cache-dir
|
||||
fi
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 2. 安装 transformers / accelerate / qwen_vl_utils
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "安装 transformers + accelerate + qwen_vl_utils ..."
|
||||
|
||||
pip install --target="$SITE_PACKAGES" --upgrade \
|
||||
"transformers>=4.40.0" \
|
||||
"accelerate>=0.30.0" \
|
||||
"qwen-vl-utils>=0.0.8" \
|
||||
"pillow>=10.0.0" \
|
||||
"safetensors>=0.4.0" \
|
||||
--no-cache-dir
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 3. 可选：安装 Flash Attention 2 (Ampere+ GPU)
|
||||
# -----------------------------------------------------------------------------
|
||||
GPU_ARCH=$(python3 -c "
|
||||
import sys, os
|
||||
sys.path.insert(0, '$SITE_PACKAGES')
|
||||
import torch
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
print(f'{major}{minor}')
|
||||
" 2>/dev/null || echo "0")
|
||||
|
||||
if [[ "$GPU_ARCH" == "8"* ]] || [[ "$GPU_ARCH" == "9"* ]]; then
|
||||
log_info "检测到 Ampere/Hopper 架构 (sm_$GPU_ARCH)，尝试安装 Flash Attention 2..."
|
||||
pip install --target="$SITE_PACKAGES" --no-build-isolation \
|
||||
"flash-attn>=2.5.0" --no-cache-dir || log_warn "Flash Attention 安装失败，将使用 eager attention"
|
||||
else
|
||||
log_warn "GPU 架构 sm_$GPU_ARCH 不支持 Flash Attention 2，跳过安装"
|
||||
fi
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 4. 验证安装
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "验证 Python 包..."
|
||||
|
||||
python3 -c "
|
||||
import sys
|
||||
sys.path.insert(0, '$SITE_PACKAGES')
|
||||
import torch
|
||||
print(f'PyTorch: {torch.__version__}')
|
||||
print(f'CUDA available: {torch.cuda.is_available()}')
|
||||
if torch.cuda.is_available():
|
||||
print(f'Device: {torch.cuda.get_device_name(0)}')
|
||||
print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB')
|
||||
"
|
||||
|
||||
python3 -c "
|
||||
import sys
|
||||
sys.path.insert(0, '$SITE_PACKAGES')
|
||||
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
|
||||
print('transformers + Qwen2.5-VL classes OK')
|
||||
"
|
||||
|
||||
python3 -c "
|
||||
import sys
|
||||
sys.path.insert(0, '$SITE_PACKAGES')
|
||||
from qwen_vl_utils import process_vision_info
|
||||
print('qwen_vl_utils OK')
|
||||
"
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 5. 预下载模型（可选，推荐）
|
||||
# -----------------------------------------------------------------------------
|
||||
VLM_MODEL="${VLM_MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
|
||||
log_info "预下载 VLM 模型: $VLM_MODEL ..."
|
||||
log_info "模型将缓存到: $VLM_MODEL_CACHE"
|
||||
|
||||
mkdir -p "$VLM_MODEL_CACHE"
|
||||
|
||||
python3 -c "
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, '$SITE_PACKAGES')
|
||||
os.environ['HF_HOME'] = '$VLM_MODEL_CACHE'
|
||||
|
||||
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
|
||||
print(f'正在下载模型: $VLM_MODEL ...')
|
||||
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
'$VLM_MODEL',
|
||||
torch_dtype='auto',
|
||||
device_map='auto',
|
||||
trust_remote_code=True,
|
||||
)
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
'$VLM_MODEL',
|
||||
trust_remote_code=True,
|
||||
)
|
||||
print('模型下载完成！')
|
||||
"
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# 6. 完成
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "安装完成！"
|
||||
log_info ""
|
||||
log_info "============================================"
|
||||
log_info " VLM 安装摘要"
|
||||
log_info "============================================"
|
||||
log_info "模型: $VLM_MODEL"
|
||||
log_info "缓存路径: $VLM_MODEL_CACHE"
|
||||
log_info "Python 包路径: $SITE_PACKAGES"
|
||||
log_info ""
|
||||
log_info "使用方法:"
|
||||
log_info " export VISION_BACKEND=vlm"
|
||||
log_info " export VLM_INTERVAL=5.0"
|
||||
log_info " ros2 run vision_yolo detector"
|
||||
log_info ""
|
||||
log_info "如需使用 AWQ 量化版（显存 3-4GB）:"
|
||||
log_info " export VLM_MODEL=Qwen/Qwen2.5-VL-3B-Instruct-AWQ"
|
||||
log_info "============================================"
|
||||
|
|
@ -1,14 +1,36 @@
|
|||
"""
|
||||
ROS2 Vision Detector — 支持 YOLOv8 CPU 与 Qwen2.5-VL GPU 双模式切换
|
||||
|
||||
运行时通过环境变量选择后端：
|
||||
VISION_BACKEND=yolo # 默认，现有 CPU YOLOv8（完全兼容原有行为）
|
||||
VISION_BACKEND=vlm # GPU Qwen2.5-VL-3B，异步推理避免阻塞 ROS
|
||||
VISION_BACKEND=hybrid # YOLO 持续检测 + VLM 定期深度分析
|
||||
|
||||
环境变量（仅 VLM/hybrid 模式有效）：
|
||||
VLM_INTERVAL — VLM 采样间隔（秒），默认 5.0
|
||||
VLM_MODEL — 模型名称，默认 Qwen/Qwen2.5-VL-3B-Instruct
|
||||
VLM_MAX_NEW_TOKENS — 最大生成 token 数，默认 256
|
||||
VLM_PROMPT — 自定义提示词（空则使用默认办公室巡检提示词）
|
||||
VLM_MIN_PIXELS — 图像最小像素数，默认 256*28*28 = 200704
|
||||
VLM_MAX_PIXELS — 图像最大像素数，默认 640*28*28 = 501760
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
from datetime import datetime
|
||||
|
||||
import rclpy
|
||||
from rclpy.node import Node
|
||||
from sensor_msgs.msg import Image
|
||||
from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose
|
||||
from nav_msgs.msg import Odometry
|
||||
from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose
|
||||
from std_msgs.msg import String
|
||||
from cv_bridge import CvBridge
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
import time
|
||||
import os
|
||||
from datetime import datetime
|
||||
import numpy as np
|
||||
|
||||
# COCO 80 类中文映射（YOLOv8n 使用）
|
||||
COCO_CN = {
|
||||
|
|
@ -43,17 +65,339 @@ def euler_from_quaternion(qx, qy, qz, qw):
|
|||
"""从四元数提取 yaw（偏航角）。"""
|
||||
siny_cosp = 2.0 * (qw * qz + qx * qy)
|
||||
cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz)
|
||||
return time.math.atan2(siny_cosp, cosy_cosp) if hasattr(time, "math") else 0.0
|
||||
import math
|
||||
return math.atan2(siny_cosp, cosy_cosp)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 后端实现
|
||||
# ============================================================================
|
||||
|
||||
class YoloBackend:
|
||||
"""YOLOv8 CPU 后端（原有逻辑封装）。"""
|
||||
|
||||
def __init__(self, node_logger):
|
||||
from ultralytics import YOLO
|
||||
self.logger = node_logger
|
||||
self.logger.info("[vision] 正在加载 YOLOv8n 模型...")
|
||||
self.model = YOLO("yolov8n.pt")
|
||||
self.logger.info("[vision] YOLOv8n 模型加载完成")
|
||||
|
||||
def detect(self, cv_img):
|
||||
"""返回 (detections_list, scene_description)。"""
|
||||
results = self.model(cv_img, verbose=False, imgsz=640)
|
||||
dets = []
|
||||
for r in results:
|
||||
for box in r.boxes:
|
||||
cls_id = int(box.cls[0])
|
||||
conf = float(box.conf[0])
|
||||
label_en = self.model.names[cls_id]
|
||||
label_cn = COCO_CN.get(label_en, label_en)
|
||||
x1, y1, x2, y2 = map(float, box.xyxy[0])
|
||||
dets.append({
|
||||
"label_en": label_en,
|
||||
"label_cn": label_cn,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
"confidence": conf,
|
||||
})
|
||||
return dets, ""
|
||||
|
||||
|
||||
class VlmBackend:
|
||||
"""Qwen2.5-VL GPU 后端：延迟加载 + 线程安全推理。"""
|
||||
|
||||
DEFAULT_PROMPT = (
|
||||
"你是一个机器人巡检助手。请分析这张办公室场景图像，识别所有可见的物体。\n"
|
||||
"\n"
|
||||
"要求：\n"
|
||||
"1. 只识别办公室常见物品：椅子、桌子、沙发、饮水机、门、窗户、柜子、电脑、"
|
||||
"绿植、垃圾桶、水杯、打印机、白板、书架、微波炉、冰箱等\n"
|
||||
"2. 忽略不相关物体（如飞机、鸟、交通灯、汽车等室外物体）\n"
|
||||
"3. 以 JSON 格式返回，包含每个物体的名称和边界框坐标\n"
|
||||
"4. 坐标使用原始图像像素值（绝对坐标），不要归一化到 0-1\n"
|
||||
"\n"
|
||||
"输出格式示例：\n"
|
||||
'{\n'
|
||||
' "objects": [\n'
|
||||
' {"label": "椅子", "bbox": [100, 200, 300, 500], "confidence": 0.92},\n'
|
||||
' {"label": "饮水机", "bbox": [400, 150, 550, 600], "confidence": 0.88}\n'
|
||||
' ],\n'
|
||||
' "scene_description": "办公室茶水间，有一台饮水机和两把椅子"\n'
|
||||
'}\n'
|
||||
"\n"
|
||||
"请只输出 JSON，不要输出其他解释文字。"
|
||||
)
|
||||
|
||||
def __init__(self, node_logger):
|
||||
self.logger = node_logger
|
||||
self.model_name = os.environ.get(
|
||||
"VLM_MODEL", "Qwen/Qwen2.5-VL-3B-Instruct"
|
||||
)
|
||||
self.max_new_tokens = int(os.environ.get("VLM_MAX_NEW_TOKENS", "256"))
|
||||
self.min_pixels = int(os.environ.get("VLM_MIN_PIXELS", str(256 * 28 * 28)))
|
||||
self.max_pixels = int(os.environ.get("VLM_MAX_PIXELS", str(640 * 28 * 28)))
|
||||
self.prompt = os.environ.get("VLM_PROMPT") or self.DEFAULT_PROMPT
|
||||
|
||||
self._model = None
|
||||
self._processor = None
|
||||
self._device = None
|
||||
self._process_vision_info = None
|
||||
self._lock = threading.Lock()
|
||||
self._load_time = 0.0
|
||||
self.infer_count = 0
|
||||
self.total_infer_time = 0.0
|
||||
|
||||
def _ensure_loaded(self):
|
||||
if self._model is not None:
|
||||
return
|
||||
with self._lock:
|
||||
if self._model is not None:
|
||||
return
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
import torch
|
||||
from transformers import (
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
AutoProcessor,
|
||||
)
|
||||
from qwen_vl_utils import process_vision_info
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
f"VLM 依赖未安装: {e}。请运行: bash scripts/install_vlm.sh"
|
||||
) from e
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
"CUDA 不可用。VLM 后端需要 GPU。"
|
||||
)
|
||||
|
||||
self._device = torch.device("cuda")
|
||||
self.logger.info(f"[VLM] 正在加载模型 {self.model_name} ...")
|
||||
|
||||
self._processor = AutoProcessor.from_pretrained(
|
||||
self.model_name,
|
||||
trust_remote_code=True,
|
||||
min_pixels=self.min_pixels,
|
||||
max_pixels=self.max_pixels,
|
||||
)
|
||||
self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype="auto",
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
self._model.eval()
|
||||
self._process_vision_info = process_vision_info
|
||||
|
||||
self._load_time = time.time() - t0
|
||||
mem_gb = torch.cuda.memory_allocated() / 1e9
|
||||
self.logger.info(
|
||||
f"[VLM] 模型加载完成，耗时 {self._load_time:.1f}s，"
|
||||
f"显存占用 {mem_gb:.2f}GB"
|
||||
)
|
||||
|
||||
def detect(self, cv_img):
|
||||
"""运行 VLM 推理。返回 (detections_list, scene_description)。"""
|
||||
self._ensure_loaded()
|
||||
t0 = time.time()
|
||||
|
||||
from PIL import Image
|
||||
|
||||
rgb_img = cv_img[:, :, ::-1]
|
||||
pil_img = Image.fromarray(rgb_img).convert("RGB")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": pil_img},
|
||||
{"type": "text", "text": self.prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
text = self._processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
image_inputs, video_inputs = self._process_vision_info(messages)
|
||||
inputs = self._processor(
|
||||
text=[text],
|
||||
images=image_inputs,
|
||||
videos=video_inputs,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
inputs = inputs.to(self._device)
|
||||
|
||||
import torch
|
||||
|
||||
with self._lock:
|
||||
with torch.no_grad():
|
||||
generated_ids = self._model.generate(
|
||||
**inputs, max_new_tokens=self.max_new_tokens
|
||||
)
|
||||
|
||||
generated_ids_trimmed = [
|
||||
out_ids[len(in_ids):]
|
||||
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
||||
]
|
||||
output_text = self._processor.batch_decode(
|
||||
generated_ids_trimmed,
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False,
|
||||
)[0]
|
||||
|
||||
elapsed = time.time() - t0
|
||||
self.infer_count += 1
|
||||
self.total_infer_time += elapsed
|
||||
self.logger.info(f"[VLM] 第 {self.infer_count} 次推理完成，耗时 {elapsed:.2f}s")
|
||||
|
||||
objects, scene_desc = self._parse_output(output_text, cv_img.shape)
|
||||
return objects, scene_desc
|
||||
|
||||
def _parse_output(self, text, img_shape):
|
||||
objects = []
|
||||
scene_desc = ""
|
||||
|
||||
# 1. 尝试 JSON 提取
|
||||
json_match = re.search(r"\{.*\}", text, re.DOTALL)
|
||||
if json_match:
|
||||
try:
|
||||
data = json.loads(json_match.group())
|
||||
for obj in data.get("objects", []):
|
||||
bbox = obj.get("bbox", [0, 0, 0, 0])
|
||||
if len(bbox) == 4:
|
||||
h, w = img_shape[:2]
|
||||
x1, y1, x2, y2 = bbox
|
||||
x1 = max(0, min(int(x1), w))
|
||||
y1 = max(0, min(int(y1), h))
|
||||
x2 = max(0, min(int(x2), w))
|
||||
y2 = max(0, min(int(y2), h))
|
||||
label = obj.get("label", "unknown")
|
||||
objects.append({
|
||||
"label_en": label,
|
||||
"label_cn": label,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
"confidence": float(obj.get("confidence", 0.85)),
|
||||
})
|
||||
scene_desc = data.get("scene_description", "")
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 2. 兜底：解析原生 Qwen grounding 格式
|
||||
if not objects:
|
||||
pattern = re.compile(
|
||||
r"<\|box_start\|>\((\d+),(\d+)\),\((\d+),(\d+)\)<\|box_end\|>"
|
||||
)
|
||||
for m in pattern.finditer(text):
|
||||
x1, y1, x2, y2 = map(int, m.groups())
|
||||
h, w = img_shape[:2]
|
||||
x1 = max(0, min(x1, w))
|
||||
y1 = max(0, min(y1, h))
|
||||
x2 = max(0, min(x2, w))
|
||||
y2 = max(0, min(y2, h))
|
||||
label = self._extract_label(text, m.start())
|
||||
objects.append({
|
||||
"label_en": label,
|
||||
"label_cn": label,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
"confidence": 0.85,
|
||||
})
|
||||
|
||||
if not scene_desc:
|
||||
scene_desc = text.split("。")[0].strip()
|
||||
if len(scene_desc) > 200:
|
||||
scene_desc = scene_desc[:200] + "..."
|
||||
|
||||
return objects, scene_desc
|
||||
|
||||
def _extract_label(self, text, box_pos):
|
||||
before = text[max(0, box_pos - 200):box_pos]
|
||||
ref_match = re.search(
|
||||
r"<\|object_ref_start\|>(.*?)<\|object_ref_end\|>", before
|
||||
)
|
||||
if ref_match:
|
||||
return ref_match.group(1).strip()
|
||||
word_match = re.search(
|
||||
r"([\u4e00-\u9fa5a-zA-Z]{1,10})\s*<\|box_start\|>", before
|
||||
)
|
||||
if word_match:
|
||||
return word_match.group(1).strip()
|
||||
return "未知物体"
|
||||
|
||||
|
||||
class HybridBackend:
|
||||
"""混合模式：YOLO 每帧检测 + VLM 定期深度分析。"""
|
||||
|
||||
def __init__(self, node_logger):
|
||||
self.logger = node_logger
|
||||
self.yolo = YoloBackend(node_logger)
|
||||
self.vlm = VlmBackend(node_logger)
|
||||
self.vlm_interval = float(os.environ.get("VLM_INTERVAL", "10.0"))
|
||||
self.last_vlm_time = 0.0
|
||||
self.last_scene_desc = ""
|
||||
self._vlm_thread = None
|
||||
self._pending_img = None
|
||||
|
||||
def detect(self, cv_img):
|
||||
# YOLO 持续运行（快）
|
||||
yolo_dets, _ = self.yolo.detect(cv_img)
|
||||
|
||||
# 周期性触发 VLM（后台线程）
|
||||
now = time.time()
|
||||
if now - self.last_vlm_time >= self.vlm_interval:
|
||||
if self._vlm_thread is None or not self._vlm_thread.is_alive():
|
||||
self._pending_img = cv_img.copy()
|
||||
self._vlm_thread = threading.Thread(
|
||||
target=self._run_vlm, daemon=True
|
||||
)
|
||||
self._vlm_thread.start()
|
||||
self.last_vlm_time = now
|
||||
|
||||
return yolo_dets, self.last_scene_desc
|
||||
|
||||
def _run_vlm(self):
|
||||
try:
|
||||
dets, desc = self.vlm.detect(self._pending_img)
|
||||
self.last_scene_desc = desc
|
||||
self.logger.info(f"[Hybrid] VLM 场景描述: {desc}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"[Hybrid] VLM 推理失败: {e}")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ROS2 节点
|
||||
# ============================================================================
|
||||
|
||||
class YoloDetector(Node):
|
||||
def __init__(self):
|
||||
super().__init__("yolo_detector")
|
||||
self.bridge = CvBridge()
|
||||
self.get_logger().info("[vision] 正在加载 YOLOv8n 模型...")
|
||||
self.model = YOLO("yolov8n.pt")
|
||||
self.get_logger().info("[vision] YOLOv8n 模型加载完成")
|
||||
|
||||
# 后端选择
|
||||
backend_type = os.environ.get("VISION_BACKEND", "yolo").lower()
|
||||
self._use_async = False
|
||||
if backend_type == "vlm":
|
||||
self.backend = VlmBackend(self.get_logger())
|
||||
self.infer_interval = float(os.environ.get("VLM_INTERVAL", "5.0"))
|
||||
self._use_async = True # VLM 慢，后台线程执行
|
||||
self.get_logger().info(
|
||||
f"[vision] VLM 模式已启用（采样间隔 {self.infer_interval}s）"
|
||||
)
|
||||
elif backend_type == "hybrid":
|
||||
self.backend = HybridBackend(self.get_logger())
|
||||
self.infer_interval = 1.0 # YOLO 节流间隔
|
||||
self._use_async = False # Hybrid 内部自己管理线程
|
||||
self.get_logger().info("[vision] Hybrid 模式已启用")
|
||||
else:
|
||||
self.backend = YoloBackend(self.get_logger())
|
||||
self.infer_interval = 1.0
|
||||
self.get_logger().info(
|
||||
"[vision] YOLO 模式（默认）。可通过 VISION_BACKEND=vlm 切换为 VLM。"
|
||||
)
|
||||
|
||||
# 订阅与发布
|
||||
self.sub = self.create_subscription(
|
||||
Image, "/camera/image_raw", self.on_image, 10)
|
||||
self.pub = self.create_publisher(
|
||||
|
|
@ -62,10 +406,12 @@ class YoloDetector(Node):
|
|||
Image, "/vision/image_annotated", 10)
|
||||
self.sub_odom = self.create_subscription(
|
||||
Odometry, "/odom", self.on_odom, 10)
|
||||
self.pub_scene = self.create_publisher(
|
||||
String, "/vision/scene_description", 10)
|
||||
|
||||
self.last_infer_time = 0.0
|
||||
self.infer_interval = 1.0
|
||||
self.frame_counter = 0
|
||||
self._async_busy = False
|
||||
|
||||
# 当前机器人位姿
|
||||
self.robot_x = 0.0
|
||||
|
|
@ -73,26 +419,18 @@ class YoloDetector(Node):
|
|||
self.robot_yaw = 0.0
|
||||
self.has_odom = False
|
||||
|
||||
# 创建输出目录：/workspace/vision_output/YYYYMMDD_HHMMSS/
|
||||
# 输出目录
|
||||
self.output_dir = os.path.join(
|
||||
"/workspace", "vision_output",
|
||||
datetime.now().strftime("%Y%m%d_%H%M%S"))
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
self.get_logger().info(f"[vision] 图像保存目录: {self.output_dir}")
|
||||
|
||||
self.get_logger().info(
|
||||
"[vision] YOLO 检测节点已启动，订阅 /camera/image_raw，"
|
||||
"推理间隔 1.0 秒，分辨率 640x480")
|
||||
|
||||
def on_odom(self, msg):
|
||||
self.robot_x = msg.pose.pose.position.x
|
||||
self.robot_y = msg.pose.pose.position.y
|
||||
q = msg.pose.pose.orientation
|
||||
# 计算 yaw
|
||||
siny_cosp = 2.0 * (q.w * q.z + q.x * q.y)
|
||||
cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
|
||||
import math
|
||||
self.robot_yaw = math.atan2(siny_cosp, cosy_cosp)
|
||||
self.robot_yaw = euler_from_quaternion(q.x, q.y, q.z, q.w)
|
||||
self.has_odom = True
|
||||
|
||||
def on_image(self, msg):
|
||||
|
|
@ -109,27 +447,48 @@ class YoloDetector(Node):
|
|||
self.get_logger().error(f"[vision] cv_bridge 转换失败: {e}")
|
||||
return
|
||||
|
||||
# 保存原始图像到时间戳目录
|
||||
raw_path = os.path.join(self.output_dir, f"frame_{self.frame_counter:04d}.jpg")
|
||||
# 保存原始图像
|
||||
raw_path = os.path.join(
|
||||
self.output_dir, f"frame_{self.frame_counter:04d}.jpg")
|
||||
cv2.imwrite(raw_path, cv_img)
|
||||
# 同时更新快速查看文件
|
||||
cv2.imwrite("/workspace/camera_latest.jpg", cv_img)
|
||||
|
||||
# YOLO 推理
|
||||
results = self.model(cv_img, verbose=False, imgsz=640)
|
||||
# 执行检测（VLM 模式使用异步避免阻塞 ROS executor）
|
||||
if self._use_async:
|
||||
if not self._async_busy:
|
||||
self._async_busy = True
|
||||
threading.Thread(
|
||||
target=self._detect_and_publish,
|
||||
args=(cv_img.copy(), msg.header),
|
||||
daemon=True,
|
||||
).start()
|
||||
else:
|
||||
self._detect_and_publish(cv_img, msg.header)
|
||||
|
||||
def _detect_and_publish(self, cv_img, header):
|
||||
base = os.path.join(
|
||||
self.output_dir, f"frame_{self.frame_counter:04d}")
|
||||
|
||||
try:
|
||||
objects, scene_desc = self.backend.detect(cv_img)
|
||||
except Exception as e:
|
||||
self.get_logger().error(f"[vision] 检测失败: {e}")
|
||||
objects = []
|
||||
scene_desc = ""
|
||||
finally:
|
||||
if self._use_async:
|
||||
self._async_busy = False
|
||||
|
||||
# 构建 ROS2 Detection2DArray
|
||||
dets = Detection2DArray()
|
||||
dets.header = msg.header
|
||||
dets.header = header
|
||||
detections_cn = []
|
||||
|
||||
detections_cn = [] # (中文标签, 置信度)
|
||||
detections_en = [] # (英文标签, 置信度)
|
||||
|
||||
for r in results:
|
||||
for box in r.boxes:
|
||||
cls_id = int(box.cls[0])
|
||||
conf = float(box.conf[0])
|
||||
label_en = self.model.names[cls_id]
|
||||
label_cn = COCO_CN.get(label_en, label_en)
|
||||
x1, y1, x2, y2 = map(float, box.xyxy[0])
|
||||
for obj in objects:
|
||||
label_en = obj["label_en"]
|
||||
label_cn = obj["label_cn"]
|
||||
conf = obj["confidence"]
|
||||
x1, y1, x2, y2 = obj["bbox"]
|
||||
|
||||
det = Detection2D()
|
||||
det.bbox.center.position.x = (x1 + x2) / 2.0
|
||||
|
|
@ -139,15 +498,14 @@ class YoloDetector(Node):
|
|||
det.bbox.size_y = y2 - y1
|
||||
|
||||
hyp = ObjectHypothesisWithPose()
|
||||
hyp.hypothesis.class_id = label_en
|
||||
hyp.hypothesis.class_id = label_en # ROS 标准用英文标识符
|
||||
hyp.hypothesis.score = conf
|
||||
det.results.append(hyp)
|
||||
dets.detections.append(det)
|
||||
|
||||
detections_en.append((label_en, conf))
|
||||
detections_cn.append((label_cn, conf))
|
||||
|
||||
# 在图像上画框和标签（中文标签）
|
||||
# 画框（中文标签）
|
||||
cv2.rectangle(cv_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||||
text = f"{label_cn} {conf:.2f}"
|
||||
cv2.putText(cv_img, text, (int(x1), int(y1) - 10),
|
||||
|
|
@ -156,17 +514,24 @@ class YoloDetector(Node):
|
|||
self.pub.publish(dets)
|
||||
|
||||
# 保存标注图像
|
||||
ann_path = os.path.join(
|
||||
self.output_dir, f"frame_{self.frame_counter:04d}_annotated.jpg")
|
||||
cv2.imwrite(ann_path, cv_img)
|
||||
cv2.imwrite(f"{base}_annotated.jpg", cv_img)
|
||||
|
||||
# 发布可视化图像
|
||||
vis_msg = self.bridge.cv2_to_imgmsg(cv_img, "bgr8")
|
||||
vis_msg.header = msg.header
|
||||
vis_msg.header = header
|
||||
self.pub_vis.publish(vis_msg)
|
||||
|
||||
# 构建中文日志
|
||||
pos_str = f"({self.robot_x:.2f}, {self.robot_y:.2f}, {self.robot_yaw:.2f}rad)" if self.has_odom else "(未知)"
|
||||
# 发布场景描述（VLM / hybrid 模式）
|
||||
if scene_desc:
|
||||
scene_msg = String()
|
||||
scene_msg.data = scene_desc
|
||||
self.pub_scene.publish(scene_msg)
|
||||
|
||||
# 日志
|
||||
pos_str = (
|
||||
f"({self.robot_x:.2f}, {self.robot_y:.2f}, {self.robot_yaw:.2f}rad)"
|
||||
if self.has_odom else "(未知)"
|
||||
)
|
||||
if detections_cn:
|
||||
items = [f"{label}({conf:.0%})" for label, conf in detections_cn]
|
||||
log_msg = (
|
||||
|
|
|
|||
Loading…
Reference in New Issue