Skip to content

Commit 0251a16

Browse files
litxaohulitxaohu
committed
[add] 新增class类别动态加载功能，添加yolo11m和yolo11s模型
1 parent 1e0e709 commit 0251a16

7 files changed

Lines changed: 150 additions & 88 deletions

File tree

‎README.md‎

Lines changed: 38 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -60,70 +60,46 @@ sudo docker pull ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest
6060

6161
**针对 RK3588:**
6262
```bash
63-
sudo docker run --rm --privileged --net=host --env DISPLAY=$DISPLAY \
64-
-v /tmp/.X11-unix:/tmp/.X11-unix \
65-
-v /dev/bus/usb:/dev/bus/usb \
63+
sudo docker run --rm --privileged --net=host \
64+
-e PYTHONUNBUFFERED=1 \
65+
-e RKNN_LOG_LEVEL=0 \
6666
--device /dev/video1:/dev/video1 \
6767
--device /dev/dri/renderD129:/dev/dri/renderD129 \
6868
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
69-
ghcr.io/litxaohu/recomputer-rk-cv/rk3588-yolo:latest
70-
python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 1
69+
ghcr.io/litxaohu/recomputer-rk-cv/rk3588-yolo:latest \
70+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1
7171
```
7272

7373
**针对 RK3576:**
7474
```bash
75-
sudo docker run --rm --privileged --net=host --env DISPLAY=$DISPLAY \
76-
-v /tmp/.X11-unix:/tmp/.X11-unix \
77-
-v /dev/bus/usb:/dev/bus/usb \
75+
sudo docker run --rm --privileged --net=host \
76+
-e PYTHONUNBUFFERED=1 \
77+
-e RKNN_LOG_LEVEL=0 \
7878
--device /dev/video0:/dev/video0 \
7979
--device /dev/dri/renderD128:/dev/dri/renderD128 \
8080
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
81-
ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest
82-
python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 0
81+
ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest \
82+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 0
8383
```
8484

85-
#### 如何预览：
86-
1. **本地显示器**：自动弹出实时检测窗口（需连接显示器并执行了 xhost）。
87-
2. **Web 浏览器**：在局域网内访问 `http://<开发板IP>:8000` 即可实时预览。
88-
89-
#### 常见问题排查：
90-
**问题：SSH 远程无屏幕运行报错 `qt.qpa.xcb: could not connect to display`**
91-
解决方案：在运行命令末尾添加 `--no_gui` 参数，强制关闭本地窗口初始化。
92-
```bash
93-
# 示例 (在原有命令末尾追加):
94-
... python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 0 --no_gui
95-
```
85+
访问方式：`http://<开发板IP>:8000`
9686

97-
### 3. 独立 Web 预览模式 (仅浏览器查看)
87+
> **注意**: 如果需要自定义类别，可以增加 `-v $(pwd)/class_config.txt:/app/class_config.txt \` 挂载和 `--class_path` 参数，程序默认使用 COCO 80 类。
9888
99-
如果您只需要通过 Web 浏览器查看预览画面（例如在远程服务器或无显示器环境下运行），可以使用专用的 Web 预览脚本：
89+
例如：
10090

101-
**针对 RK3588:**
10291
```bash
10392
sudo docker run --rm --privileged --net=host \
10493
-e PYTHONUNBUFFERED=1 \
10594
-e RKNN_LOG_LEVEL=0 \
95+
-v $(pwd)/class_config.txt:/app/class_config.txt \
10696
--device /dev/video1:/dev/video1 \
10797
--device /dev/dri/renderD129:/dev/dri/renderD129 \
10898
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
10999
ghcr.io/litxaohu/recomputer-rk-cv/rk3588-yolo:latest \
110-
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1
111-
```
112-
113-
**针对 RK3576:**
114-
```bash
115-
sudo docker run --rm --privileged --net=host \
116-
-e PYTHONUNBUFFERED=1 \
117-
-e RKNN_LOG_LEVEL=0 \
118-
--device /dev/video0:/dev/video0 \
119-
--device /dev/dri/renderD128:/dev/dri/renderD128 \
120-
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
121-
ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest \
122-
python web_detection.py --model_path model/yolo11n.rknn --camera_id 0
100+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1 --class_path class_config.txt
123101
```
124102

125-
访问方式：`http://<开发板IP>:8000`
126-
127103
---
128104

129105
## 🔌 API 接口文档
@@ -190,7 +166,26 @@ curl -X POST "http://127.0.0.1:8000/api/models/yolo11/predict"
190166
- **请求体 (JSON):** `{"obj_thresh": 0.3, "nms_thresh": 0.5}`
191167
- **响应:** `{"status": "success"}`
192168

193-
### 3. 实时视频流接口 (Video Feed)
169+
### 3. 命令行参数说明
170+
171+
运行 `web_detection.py` 时支持以下参数：
172+
173+
| 参数 | 说明 | 默认值 |
174+
| :--- | :--- | :--- |
175+
| `--model_path` | RKNN 模型文件路径 | (必填) |
176+
| `--camera_id` | 摄像头设备 ID (如 `/dev/video1` 则填 1) | 1 |
177+
| `--video_path` | 视频文件路径 (若提供则忽略 camera_id) | None |
178+
| `--class_path` | 自定义类别配置文件路径 (class_config.txt) | None (默认 COCO 80类) |
179+
| `--host` | Web 服务器监听地址 | `0.0.0.0` |
180+
| `--port` | Web 服务器端口 | 8000 |
181+
182+
#### 自定义类别配置 (class_config.txt) 格式：
183+
使用双引号命名分类，不同分类之间使用逗号隔开，例如：
184+
`"person", "bicycle", "car", "motorbike"`
185+
186+
---
187+
188+
## 实时视频流接口 (Video Feed)
194189

195190
获取带有检测框绘制的实时 MJPEG 视频流，可直接嵌入 HTML `<img>` 标签。
196191

@@ -211,12 +206,13 @@ curl -X POST "http://127.0.0.1:8000/api/models/yolo11/predict"
211206

212207
---
213208

214-
## �️ 开发者指南 (量产建议)
209+
## 🛠️ 开发者指南 (量产建议)
215210
### 代码说明
216-
- `realtime_detection.py`:
211+
- `web_detection.py`:
217212
- **双模支持**: 集成 FastAPI，同时支持本地渲染和 MJPEG 流式输出。
218213
- **环境自适应**: 自动检测 `DISPLAY` 环境变量，无环境时静默跳过 GUI 初始化。
219214
- **RKNN 推理**: 封装了 RKNN 初始化、加载模型、多核推理逻辑。
215+
- **动态加载**: 支持通过 `--class_path` 动态加载类别配置。
220216
- **后处理**: YOLOv11 专用的 Box 解码与 NMS 逻辑。
221217

222218
### 修改模型

‎src/rk3576/README.md‎

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
## 目录结构
1111
- `lib/`：包含 RK3576 版 `librknnrt.so`。
1212
- `model/`：存放针对 RK3576 转换的 `.rknn` 模型。
13-
- `realtime_detection.py`：主程序。
13+
- `web_detection.py`：主程序（支持 Web 预览与 API）。
1414

1515
## 快速开始
1616

@@ -26,43 +26,34 @@ xhost +local:docker
2626

2727
#### 步骤 B：一键运行
2828
```bash
29-
sudo docker run --rm --privileged --net=host --env DISPLAY=$DISPLAY \
30-
-v /tmp/.X11-unix:/tmp/.X11-unix \
31-
-v /dev/bus/usb:/dev/bus/usb \
29+
sudo docker run --rm --privileged --net=host \
30+
-e PYTHONUNBUFFERED=1 \
31+
-e RKNN_LOG_LEVEL=0 \
3232
--device /dev/video0:/dev/video0 \
3333
--device /dev/dri/renderD128:/dev/dri/renderD128 \
3434
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
3535
ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest \
36-
python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 0
36+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 0
3737
```
38+
访问方式：`http://<开发板IP>:8000`
3839

39-
#### 如何预览：
40-
1. **本地显示器**：自动弹出实时检测窗口（需连接显示器并执行了 xhost）。
41-
2. **Web 浏览器**：在局域网内访问 `http://<开发板IP>:8000` 即可实时预览。
4240

43-
#### 常见问题排查：
44-
**问题：SSH 远程无屏幕运行报错 `qt.qpa.xcb: could not connect to display`**
45-
解决方案：在运行命令末尾添加 `--no_gui` 参数，强制关闭本地窗口初始化。
46-
```bash
47-
# 示例 (在原有命令末尾追加):
48-
... python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 0 --no_gui
49-
```
50-
51-
### 2. 独立 Web 预览模式 (仅浏览器查看)
41+
> **注意**: 如果需要自定义类别，可以增加 `-v $(pwd)/class_config.txt:/app/class_config.txt \` 挂载和 `--class_path` 参数，程序默认使用 COCO 80 类。
5242
53-
如果您只需要通过 Web 浏览器查看预览画面（例如在远程服务器或无显示器环境下运行），可以使用专用的 Web 预览脚本：
43+
例如：
5444

5545
```bash
5646
sudo docker run --rm --privileged --net=host \
5747
-e PYTHONUNBUFFERED=1 \
5848
-e RKNN_LOG_LEVEL=0 \
49+
-v $(pwd)/class_config.txt:/app/class_config.txt \
5950
--device /dev/video0:/dev/video0 \
6051
--device /dev/dri/renderD128:/dev/dri/renderD128 \
6152
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
6253
ghcr.io/litxaohu/recomputer-rk-cv/rk3576-yolo:latest \
63-
python web_detection.py --model_path model/yolo11n.rknn --camera_id 0
54+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 0 --class_path class_config.txt
6455
```
65-
访问方式：`http://<开发板IP>:8000`
56+
6657

6758
---
6859

@@ -139,12 +130,13 @@ curl -X POST "http://127.0.0.1:8000/api/models/yolo11/predict"
139130

140131
---
141132

142-
## ️ 开发者指南 (量产建议)
133+
## 🛠️ 开发者指南 (量产建议)
143134
### 代码说明
144-
- `realtime_detection.py`:
135+
- `web_detection.py`:
145136
- **双模支持**: 集成 FastAPI，同时支持本地渲染和 MJPEG 流式输出。
146137
- **环境自适应**: 自动检测 `DISPLAY` 环境变量，无环境时静默跳过 GUI 初始化。
147138
- **RKNN 推理**: 封装了 RKNN 初始化、加载模型、多核推理逻辑。
139+
- **动态加载**: 支持通过 `--class_path` 动态加载类别配置。
148140
- **后处理**: YOLOv11 专用的 Box 解码与 NMS 逻辑。
149141

150142
### 修改模型

‎src/rk3576/model/yolo11m.rknn‎

25.8 MB
Binary file not shown.

‎src/rk3576/model/yolo11s.rknn‎

13.8 MB
Binary file not shown.

‎src/rk3576/web_detection.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,50 @@
2626
NMS_THRESH = 0.45
2727
IMG_SIZE = (640, 640) # (width, height)
2828

29-
CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
29+
# 默认类别定义 (COCO 80类)
30+
DEFAULT_CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
3031
"fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant",
3132
"bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
3233
"baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ",
3334
"spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa",
3435
"pottedplant","bed","diningtable","toilet ","tvmonitor","laptop ","mouse ","remote ","keyboard ","cell phone","microwave ",
3536
"oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush ")
3637

38+
CLASSES = DEFAULT_CLASSES
39+
40+
def load_classes(path):
41+
"""
42+
从文件加载类别，支持双引号和逗号分隔的格式
43+
例如: "person", "bicycle", "car"
44+
"""
45+
global CLASSES
46+
if not path or not os.path.exists(path):
47+
CLASSES = DEFAULT_CLASSES
48+
return
49+
50+
try:
51+
with open(path, 'r', encoding='utf-8') as f:
52+
content = f.read().strip()
53+
# 简单的解析逻辑：移除换行，按逗号分割，去除空格和双引号
54+
import re
55+
# 匹配双引号内的内容
56+
items = re.findall(r'"([^"]*)"', content)
57+
if items:
58+
CLASSES = tuple(items)
59+
print(f"Successfully loaded {len(CLASSES)} classes from {path}")
60+
else:
61+
# 备选方案：如果没匹配到双引号，尝试按逗号分割
62+
items = [item.strip().strip('"') for item in content.split(',') if item.strip()]
63+
if items:
64+
CLASSES = tuple(items)
65+
print(f"Loaded {len(CLASSES)} classes from {path} (fallback parsing)")
66+
else:
67+
print(f"Warning: No classes found in {path}, using default COCO classes")
68+
CLASSES = DEFAULT_CLASSES
69+
except Exception as e:
70+
print(f"Error loading classes from {path}: {e}. Using default COCO classes")
71+
CLASSES = DEFAULT_CLASSES
72+
3773
# 动态配置参数
3874
class DetectionConfig:
3975
def __init__(self):
@@ -481,6 +517,7 @@ def main():
481517
parser.add_argument('--model_path', type=str, required=True, help='RKNN model path')
482518
parser.add_argument('--camera_id', type=int, default=1, help='Camera device ID (default: 1 for /dev/video1)')
483519
parser.add_argument('--video_path', type=str, help='Path to video file (overrides camera_id)')
520+
parser.add_argument('--class_path', type=str, help='Path to class_config.txt file for dynamic category loading')
484521
parser.add_argument('--host', type=str, default='0.0.0.0', help='Web server host')
485522
parser.add_argument('--port', type=int, default=8000, help='Web server port')
486523
args = parser.parse_args()
@@ -489,6 +526,10 @@ def main():
489526
print("Error: RKNN-Toolkit-Lite2 is not available.")
490527
return
491528

529+
# 加载自定义类别
530+
if args.class_path:
531+
load_classes(args.class_path)
532+
492533
# 启动 Web 服务器线程
493534
web_thread = threading.Thread(target=run_fastapi, args=(args.host, args.port), daemon=True)
494535
web_thread.start()

‎src/rk3588/README.md‎

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
## 目录结构
1111
- `lib/`：包含 RK3588 版 `librknnrt.so`。
1212
- `model/`：存放针对 RK3588 转换的 `.rknn` 模型。
13-
- `realtime_detection.py`：主程序。
13+
- `web_detection.py`：主程序（支持 Web 预览与 API）。
1414

1515
## 快速开始
1616

@@ -25,44 +25,35 @@ xhost +local:docker
2525
```
2626

2727
#### 步骤 B：一键运行
28+
2829
```bash
29-
sudo docker run --rm --privileged --net=host --env DISPLAY=$DISPLAY \
30-
-v /tmp/.X11-unix:/tmp/.X11-unix \
31-
-v /dev/bus/usb:/dev/bus/usb \
30+
sudo docker run --rm --privileged --net=host \
31+
-e PYTHONUNBUFFERED=1 \
32+
-e RKNN_LOG_LEVEL=0 \
3233
--device /dev/video1:/dev/video1 \
3334
--device /dev/dri/renderD129:/dev/dri/renderD129 \
3435
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
3536
ghcr.io/litxaohu/recomputer-rk-cv/rk3588-yolo:latest \
36-
python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 1
37+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1
3738
```
39+
访问方式：`http://<开发板IP>:8000`
3840

39-
#### 如何预览：
40-
1. **本地显示器**：自动弹出实时检测窗口（需连接显示器并执行了 xhost）。
41-
2. **Web 浏览器**：在局域网内访问 `http://<开发板IP>:8000` 即可实时预览。
42-
43-
#### 常见问题排查：
44-
**问题：SSH 远程无屏幕运行报错 `qt.qpa.xcb: could not connect to display`**
45-
解决方案：在运行命令末尾添加 `--no_gui` 参数，强制关闭本地窗口初始化。
46-
```bash
47-
# 示例 (在原有命令末尾追加):
48-
... python realtime_detection.py --model_path model/yolo11n.rknn --camera_id 1 --no_gui
49-
```
5041

51-
### 2. 独立 Web 预览模式 (仅浏览器查看)
42+
> **注意**: 如果需要自定义类别，可以增加 `-v $(pwd)/class_config.txt:/app/class_config.txt \` 挂载和 `--class_path` 参数，程序默认使用 COCO 80 类。
5243
53-
如果您只需要通过 Web 浏览器查看预览画面（例如在远程服务器或无显示器环境下运行），可以使用专用的 Web 预览脚本：
44+
例如：
5445

5546
```bash
5647
sudo docker run --rm --privileged --net=host \
5748
-e PYTHONUNBUFFERED=1 \
5849
-e RKNN_LOG_LEVEL=0 \
50+
-v $(pwd)/class_config.txt:/app/class_config.txt \
5951
--device /dev/video1:/dev/video1 \
6052
--device /dev/dri/renderD129:/dev/dri/renderD129 \
6153
-v /proc/device-tree/compatible:/proc/device-tree/compatible \
6254
ghcr.io/litxaohu/recomputer-rk-cv/rk3588-yolo:latest \
63-
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1
55+
python web_detection.py --model_path model/yolo11n.rknn --camera_id 1 --class_path class_config.txt
6456
```
65-
访问方式：`http://<开发板IP>:8000`
6657

6758
---
6859

@@ -139,12 +130,13 @@ curl -X POST "http://127.0.0.1:8000/api/models/yolo11/predict"
139130

140131
---
141132

142-
## ️ 开发者指南 (量产建议)
133+
## 🛠️ 开发者指南 (量产建议)
143134
### 代码说明
144-
- `realtime_detection.py`:
135+
- `web_detection.py`:
145136
- **双模支持**: 集成 FastAPI，同时支持本地渲染和 MJPEG 流式输出。
146137
- **环境自适应**: 自动检测 `DISPLAY` 环境变量，无环境时静默跳过 GUI 初始化。
147138
- **RKNN 推理**: 封装了 RKNN 初始化、加载模型、多核推理逻辑。
139+
- **动态加载**: 支持通过 `--class_path` 动态加载类别配置。
148140
- **后处理**: YOLOv11 专用的 Box 解码与 NMS 逻辑。
149141

150142
### 修改模型

0 commit comments

Comments
 (0)