From 4a2138b8f473a3ceae55e7bc46d0bd41e698a85e Mon Sep 17 00:00:00 2001 From: inteintegrity Date: Mon, 27 Jul 2026 16:46:19 +0800 Subject: [PATCH] feat: add RK3576 airborne object detection (YOLOv8s, VisDrone) package Add a new RK3576 model package for airborne/aerial object detection based on YOLOv8s trained on the VisDrone dataset (11 aerial classes). - src/rk3576_yolov8_airborne/: full package (model, lib, py_utils, video, web_detection.py, README in EN/ZH) following the existing RK-CV layout - docker/rk3576/yolov8_airborne.dockerfile: build config mirroring yolov8 - .github/workflows/docker-build.yml: add rk3576_yolov8_airborne to the build matrix and push path triggers - README.md / README_zh.md: list the new image in the pull commands --- .github/workflows/docker-build.yml | 2 + README.md | 1 + README_zh.md | 1 + docker/rk3576/yolov8_airborne.dockerfile | 53 ++ src/rk3576_yolov8_airborne/README.md | 163 ++++ src/rk3576_yolov8_airborne/README_zh.md | 163 ++++ src/rk3576_yolov8_airborne/lib/librknnrt.so | 3 + .../model/yolov8s_airborne.rknn | 3 + .../py_utils/__init__.py | 0 .../py_utils/coco_utils.py | 176 ++++ src/rk3576_yolov8_airborne/requirements.txt | 5 + ...nux_2_17_aarch64.manylinux2014_aarch64.whl | 3 + src/rk3576_yolov8_airborne/video/test.mp4 | 3 + src/rk3576_yolov8_airborne/web_detection.py | 864 ++++++++++++++++++ 14 files changed, 1440 insertions(+) create mode 100644 docker/rk3576/yolov8_airborne.dockerfile create mode 100644 src/rk3576_yolov8_airborne/README.md create mode 100644 src/rk3576_yolov8_airborne/README_zh.md create mode 100644 src/rk3576_yolov8_airborne/lib/librknnrt.so create mode 100644 src/rk3576_yolov8_airborne/model/yolov8s_airborne.rknn create mode 100644 src/rk3576_yolov8_airborne/py_utils/__init__.py create mode 100644 src/rk3576_yolov8_airborne/py_utils/coco_utils.py create mode 100644 src/rk3576_yolov8_airborne/requirements.txt create mode 100644 src/rk3576_yolov8_airborne/rknn-toolkit-lite2-packages/rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl create mode 100644 src/rk3576_yolov8_airborne/video/test.mp4 create mode 100644 src/rk3576_yolov8_airborne/web_detection.py diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index 296840d..3f9c8a1 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -18,6 +18,7 @@ on: - 'src/rk3576_yolov8_pose/**' - 'src/rk3588_yolov8_seg/**' - 'src/rk3576_yolov8_seg/**' + - 'src/rk3576_yolov8_airborne/**' - 'src/rk3588_whisper/**' - 'src/rk3576_whisper/**' - 'docker/**' @@ -58,6 +59,7 @@ jobs: - rk3576_yolov8_pose - rk3588_yolov8_seg - rk3576_yolov8_seg + - rk3576_yolov8_airborne - rk3588_whisper - rk3576_whisper diff --git a/README.md b/README.md index d4c579b..1347f19 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,7 @@ sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3588-yolo11:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolo11:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3588-yolov8-obb:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8-obb:latest +sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest ``` #### Step C: Run with One Click diff --git a/README_zh.md b/README_zh.md index 160d9aa..8bcc800 100644 --- a/README_zh.md +++ b/README_zh.md @@ -58,6 +58,7 @@ sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3588-yolo11:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolo11:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3588-yolov8-obb:latest sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8-obb:latest +sudo docker pull ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest ``` #### 步骤 C：一键运行 diff --git a/docker/rk3576/yolov8_airborne.dockerfile b/docker/rk3576/yolov8_airborne.dockerfile new file mode 100644 index 0000000..ffb480d --- /dev/null +++ b/docker/rk3576/yolov8_airborne.dockerfile @@ -0,0 +1,53 @@ +# Use Python 3.11 slim image as base (matching the wheel requirement) +FROM python:3.11-slim + +# Set working directory +WORKDIR /app + +# Install system dependencies for OpenCV and others +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 \ + libgomp1 \ + libsm6 \ + libxext6 \ + libxrender1 \ + libxcb-cursor0 \ + libxcb-xinerama0 \ + libxcb-keysyms1 \ + libxcb-image0 \ + libxcb-shm0 \ + libxcb-icccm4 \ + libxcb-sync1 \ + libxcb-xfixes0 \ + libxcb-shape0 \ + libxcb-randr0 \ + libxcb-render-util0 \ + libxkbcommon-x11-0 \ + && rm -rf /var/lib/apt/lists/* + +# Copy requirements file +COPY requirements.txt . + +# Install Python dependencies +RUN pip install --no-cache-dir -r requirements.txt + +# Copy the RKNN Toolkit Lite2 wheel +COPY rknn-toolkit-lite2-packages/rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl . + +# Install the local wheel +RUN pip install --no-cache-dir rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl + +# Copy librknnrt.so to /usr/lib/ +COPY lib/librknnrt.so /usr/lib/ +RUN chmod 755 /usr/lib/librknnrt.so + +# Copy the rest of the application code +COPY . . + +# Expose port for Web Preview +EXPOSE 8000 + +# Set the default command to run the airborne detection script +# This script now supports both GUI (if DISPLAY is available) and Web Preview (on port 8000) +CMD ["python", "web_detection.py", "--model_path", "model/yolov8s_airborne.rknn", "--video_path", "video/test.mp4"] diff --git a/src/rk3576_yolov8_airborne/README.md b/src/rk3576_yolov8_airborne/README.md new file mode 100644 index 0000000..244405d --- /dev/null +++ b/src/rk3576_yolov8_airborne/README.md @@ -0,0 +1,163 @@ +# RK3576 Airborne Object Detection (YOLOv8s) Deployment Guide + +[English] | [中文](./README_zh.md) + +This directory contains a YOLOv8s airborne object detection inference package optimized for RK3576, trained on the VisDrone dataset for drone/aerial view object detection. + +## Core Features +- **Hardware Acceleration**: Optimized for RK3576's 6 TOPS NPU architecture. +- **Latest Driver**: Integrates the 5th generation NPU runtime library supporting RK3576. +- **Flexible Input**: Supports camera and local MP4 video input. +- **VisDrone Classes**: Pre-configured with 11 VisDrone aerial categories (pedestrian, people, bicycle, car, van, truck, tricycle, awning-tricycle, bus, motor, others). + +## Directory Structure +- `lib/`: Contains `librknnrt.so` for RK3576. +- `model/`: Stores the `.rknn` model converted for RK3576 (e.g., `yolov8s_airborne.rknn`). +- `py_utils/`: Inference engine wrapper and post-processing utilities (letterbox + coordinate restoration). +- `web_detection.py`: Main program (supports Web preview and API). +- `video/`: Sample test video. + +## VisDrone Classes (11) + +| ID | Class | +|----|-------| +| 0 | pedestrian | +| 1 | people | +| 2 | bicycle | +| 3 | car | +| 4 | van | +| 5 | truck | +| 6 | tricycle | +| 7 | awning-tricycle | +| 8 | bus | +| 9 | motor | +| 10 | others | + +## Quick Start + +### 1. Run the Project (One command, dual-mode preview) + +This project supports simultaneous preview via **Local GUI** and **Web Browser**. The program automatically detects the display environment and downgrades to Web mode if no display is connected. + +#### Step A: Configure Display Permissions (Optional) +If you have a monitor connected and want to see the window locally: +```bash +xhost +local:docker +``` + +#### Step B: One-click Run +```bash +sudo docker run --rm --privileged --net=host \ + -e PYTHONUNBUFFERED=1 \ + -e RKNN_LOG_LEVEL=0 \ + --device /dev/video0:/dev/video0 \ + --device /dev/dri/renderD128:/dev/dri/renderD128 \ + -v /proc/device-tree/compatible:/proc/device-tree/compatible \ + ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest \ + python3 web_detection.py --model_path model/yolov8s_airborne.rknn --video video/test.mp4 +``` +Access via: `http://:8000` + +> **Note**: The program defaults to the 11 VisDrone classes. If you need custom classes, you can add `-v $(pwd)/class_config.txt:/app/class_config.txt \` mount and the `--class_path` parameter. + +Example: + +```bash +sudo docker run --rm --privileged --net=host \ + -e PYTHONUNBUFFERED=1 \ + -e RKNN_LOG_LEVEL=0 \ + -v $(pwd)/class_config.txt:/app/class_config.txt \ + --device /dev/video0:/dev/video0 \ + --device /dev/dri/renderD128:/dev/dri/renderD128 \ + -v /proc/device-tree/compatible:/proc/device-tree/compatible \ + ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest \ + python3 web_detection.py --model_path model/yolov8s_airborne.rknn --video video/test.mp4 --class_path class_config.txt +``` + +--- + +## 🔌 API Documentation + +This project provides RESTful interfaces compatible with the Ultralytics Cloud API standard, supporting object detection via image, video uploads or direct camera calls. + +### 1. Model Inference Interface (Predict) + +**Endpoint:** `POST /api/models/yolov8/predict` + +#### Request Parameters (Multipart/Form-Data): +- `file`: (Optional) Image file to be detected. +- `video`: (Optional) MP4 video file to be detected. +- `timestamp`: (Optional) Timestamp in the video file (seconds), returns detection results for the frame at that point. Default is 0. +- `realtime`: (Optional) Boolean. If `true` or if no `file`/`video` parameters are provided, returns detection results for the current camera frame. +- `conf`: (Optional) Confidence threshold for a single request, range 0.0-1.0. +- `iou`: (Optional) NMS IOU threshold for a single request, range 0.0-1.0. + +#### Usage Examples: + +**1. Image Detection:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "file=@/home/cat/001.jpg" +``` + +**2. Video Specific Frame Detection:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "video=@/home/cat/test.mp4" -F "timestamp=5.5" +``` + +**3. Get Current Camera Frame Detection:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "realtime=true" +# Or without file parameters +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" +``` + +#### Response Format (JSON): +```json +{ + "success": true, + "source": "video frame at 5.5s", + "predictions": [ + { + "class": "pedestrian", + "confidence": 0.92, + "box": { "x1": 100, "y1": 200, "x2": 300, "y2": 500 } + } + ], + "image": { "width": 1280, "height": 720 } +} +``` + +### 2. System Configuration Interface (Config) + +Used to dynamically adjust thresholds for real-time video streams and default inference. + +#### Get Current Configuration +- **Endpoint:** `GET /api/config` +- **Response:** `{"obj_thresh": 0.25, "nms_thresh": 0.45}` + +#### Update System Configuration +- **Endpoint:** `POST /api/config` +- **Request Body (JSON):** `{"obj_thresh": 0.3, "nms_thresh": 0.5}` +- **Response:** `{"status": "success"}` + +### 3. Real-time Video Stream Interface (Video Feed) + +Get real-time MJPEG video stream with detection boxes drawn, can be directly embedded in HTML `` tags. + +- **Endpoint:** `GET /api/video_feed` +- **Example Usage:** `` + +--- + +## 🛠️ Developer Guide (Production Recommendations) +### Code Description +- `web_detection.py`: + - **Dual-mode Support**: Integrates FastAPI, supporting both local rendering and MJPEG streaming output. + - **Environment Adaptive**: Automatically detects the `DISPLAY` environment variable, silently skipping GUI initialization if not present. + - **RKNN Inference**: Encapsulates RKNN initialization, model loading, and multi-core inference logic. + - **Dynamic Loading**: Supports dynamic class configuration loading via `--class_path`. + - **Post-processing**: YOLOv8 specific Box decoding (DFL) and NMS logic. + +### Modifying Models +1. Place the trained and converted .rknn model into the `model/` directory. +2. Add the `--model_path` argument to the running command to point to the new model. diff --git a/src/rk3576_yolov8_airborne/README_zh.md b/src/rk3576_yolov8_airborne/README_zh.md new file mode 100644 index 0000000..0a8bb5c --- /dev/null +++ b/src/rk3576_yolov8_airborne/README_zh.md @@ -0,0 +1,163 @@ +# RK3576 机载目标检测 (YOLOv8s) 部署指南 + +[English] | [中文](./README_zh.md) + +本目录包含针对 RK3576 优化的 YOLOv8s 机载目标检测推理包，基于 VisDrone 数据集训练，适用于无人机/航拍视角的目标检测。 + +## 核心特性 +- **硬件加速**：针对 RK3576 的 6 TOPS NPU 架构进行了优化。 +- **最新驱动**：集成支持 RK3576 的第 5 代 NPU 运行时库。 +- **灵活输入**：支持摄像头和本地 MP4 视频输入。 +- **VisDrone 类别**：预置 11 个 VisDrone 航拍类别 (pedestrian, people, bicycle, car, van, truck, tricycle, awning-tricycle, bus, motor, others)。 + +## 目录结构 +- `lib/`：包含 RK3576 版 `librknnrt.so`。 +- `model/`：存放针对 RK3576 转换的 `.rknn` 模型 (如 `yolov8s_airborne.rknn`)。 +- `py_utils/`：推理引擎封装与后处理工具 (letterbox + 坐标还原)。 +- `web_detection.py`：主程序（支持 Web 预览与 API）。 +- `video/`：示例测试视频。 + +## VisDrone 类别 (11 类) + +| ID | 类别 | +|----|-------| +| 0 | pedestrian | +| 1 | people | +| 2 | bicycle | +| 3 | car | +| 4 | van | +| 5 | truck | +| 6 | tricycle | +| 7 | awning-tricycle | +| 8 | bus | +| 9 | motor | +| 10 | others | + +## 快速开始 + +### 1. 运行项目 (一条命令，双模预览) + +本项目支持 **本地 GUI** 与 **Web 浏览器** 双模式同时预览。程序会自动检测显示器环境，无显示器时自动降级为 Web 模式。 + +#### 步骤 A：配置显示权限 (可选) +如果您连接了显示器并希望在本地看到窗口： +```bash +xhost +local:docker +``` + +#### 步骤 B：一键运行 +```bash +sudo docker run --rm --privileged --net=host \ + -e PYTHONUNBUFFERED=1 \ + -e RKNN_LOG_LEVEL=0 \ + --device /dev/video0:/dev/video0 \ + --device /dev/dri/renderD128:/dev/dri/renderD128 \ + -v /proc/device-tree/compatible:/proc/device-tree/compatible \ + ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest \ + python3 web_detection.py --model_path model/yolov8s_airborne.rknn --video video/test.mp4 +``` +访问方式：`http://<开发板IP>:8000` + +> **注意**: 程序默认使用 11 个 VisDrone 类别。如果需要自定义类别，可以增加 `-v $(pwd)/class_config.txt:/app/class_config.txt \` 挂载和 `--class_path` 参数。 + +例如： + +```bash +sudo docker run --rm --privileged --net=host \ + -e PYTHONUNBUFFERED=1 \ + -e RKNN_LOG_LEVEL=0 \ + -v $(pwd)/class_config.txt:/app/class_config.txt \ + --device /dev/video0:/dev/video0 \ + --device /dev/dri/renderD128:/dev/dri/renderD128 \ + -v /proc/device-tree/compatible:/proc/device-tree/compatible \ + ghcr.io/seeed-projects/recomputer-rk-cv/rk3576-yolov8_airborne:latest \ + python3 web_detection.py --model_path model/yolov8s_airborne.rknn --video video/test.mp4 --class_path class_config.txt +``` + +--- + +## 🔌 API 接口文档 + +本项目提供了兼容 Ultralytics Cloud API 标准的 RESTful 接口，支持通过 HTTP POST 请求上传图片、视频或直接调用摄像头进行目标检测。 + +### 1. 模型推理接口 (Predict) + +**Endpoint:** `POST /api/models/yolov8/predict` + +#### 请求参数 (Multipart/Form-Data): +- `file`: (可选) 待检测的图片文件。 +- `video`: (可选) 待检测的 MP4 视频文件。 +- `timestamp`: (可选) 视频文件的时间戳（单位：秒），返回该时间点的视频帧检测结果。默认为 0。 +- `realtime`: (可选) 布尔值。若为 `true` 或未提供 `file`/`video` 参数，则返回摄像头当前帧的检测结果。 +- `conf`: (可选) 单次请求的置信度阈值，范围 0.0-1.0。 +- `iou`: (可选) 单次请求的 NMS IOU 阈值，范围 0.0-1.0。 + +#### 调用示例: + +**1. 图片检测:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "file=@/home/cat/001.jpg" +``` + +**2. 视频特定时间帧检测:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "video=@/home/cat/test.mp4" -F "timestamp=5.5" +``` + +**3. 获取摄像头当前帧检测:** +```bash +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" -F "realtime=true" +# 或者不传文件参数 +curl -X POST "http://127.0.0.1:8000/api/models/yolov8/predict" +``` + +#### 响应格式 (JSON): +```json +{ + "success": true, + "source": "video frame at 5.5s", + "predictions": [ + { + "class": "pedestrian", + "confidence": 0.92, + "box": { "x1": 100, "y1": 200, "x2": 300, "y2": 500 } + } + ], + "image": { "width": 1280, "height": 720 } +} +``` + +### 2. 系统配置接口 (Config) + +用于动态调整实时视频流和默认推理的阈值。 + +#### 获取当前配置 +- **Endpoint:** `GET /api/config` +- **响应:** `{"obj_thresh": 0.25, "nms_thresh": 0.45}` + +#### 更新系统配置 +- **Endpoint:** `POST /api/config` +- **请求体 (JSON):** `{"obj_thresh": 0.3, "nms_thresh": 0.5}` +- **响应:** `{"status": "success"}` + +### 3. 实时视频流接口 (Video Feed) + +获取带有检测框绘制的实时 MJPEG 视频流，可直接嵌入 HTML `` 标签。 + +- **Endpoint:** `GET /api/video_feed` +- **使用示例:** `` + +--- + +## 🛠️ 开发者指南 (量产建议) +### 代码说明 +- `web_detection.py`: + - **双模支持**: 集成 FastAPI，同时支持本地渲染和 MJPEG 流式输出。 + - **环境自适应**: 自动检测 `DISPLAY` 环境变量，无环境时静默跳过 GUI 初始化。 + - **RKNN 推理**: 封装了 RKNN 初始化、加载模型、多核推理逻辑。 + - **动态加载**: 支持通过 `--class_path` 动态加载类别配置。 + - **后处理**: YOLOv8 专用的 Box 解码 (DFL) 与 NMS 逻辑。 + +### 修改模型 +1. 将训练好并转换完成的 .rknn 模型放入 `model/` 目录。 +2. 运行命令时可添加 `--model_path` 参数指向新模型。 diff --git a/src/rk3576_yolov8_airborne/lib/librknnrt.so b/src/rk3576_yolov8_airborne/lib/librknnrt.so new file mode 100644 index 0000000..32e0ddd --- /dev/null +++ b/src/rk3576_yolov8_airborne/lib/librknnrt.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cf6ea624489225ddc9d334b370c64083b0fd9d6a5c9166eed47d999cc85a806b +size 7258952 diff --git a/src/rk3576_yolov8_airborne/model/yolov8s_airborne.rknn b/src/rk3576_yolov8_airborne/model/yolov8s_airborne.rknn new file mode 100644 index 0000000..8dc591d --- /dev/null +++ b/src/rk3576_yolov8_airborne/model/yolov8s_airborne.rknn @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40de56c23df646eecfb982841b82aaa2cb2317f9be7a67a978f164e7cdebf59b +size 29750134 diff --git a/src/rk3576_yolov8_airborne/py_utils/__init__.py b/src/rk3576_yolov8_airborne/py_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/rk3576_yolov8_airborne/py_utils/coco_utils.py b/src/rk3576_yolov8_airborne/py_utils/coco_utils.py new file mode 100644 index 0000000..713257c --- /dev/null +++ b/src/rk3576_yolov8_airborne/py_utils/coco_utils.py @@ -0,0 +1,176 @@ +from copy import copy +import os +import cv2 +import numpy as np +import json + +class Letter_Box_Info(): + def __init__(self, shape, new_shape, w_ratio, h_ratio, dw, dh, pad_color) -> None: + self.origin_shape = shape + self.new_shape = new_shape + self.w_ratio = w_ratio + self.h_ratio = h_ratio + self.dw = dw + self.dh = dh + self.pad_color = pad_color + + +def coco_eval_with_json(anno_json, pred_json): + from pycocotools.coco import COCO + from pycocotools.cocoeval import COCOeval + anno = COCO(anno_json) + pred = anno.loadRes(pred_json) + eval = COCOeval(anno, pred, 'bbox') + # eval.params.useCats = 0 + # eval.params.maxDets = list((100, 300, 1000)) + # a = np.array(list(range(50, 96, 1)))/100 + # eval.params.iouThrs = a + eval.evaluate() + eval.accumulate() + eval.summarize() + map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) + + print('map --> ', map) + print('map50--> ', map50) + print('map75--> ', eval.stats[2]) + print('map85--> ', eval.stats[-2]) + print('map95--> ', eval.stats[-1]) + +class COCO_test_helper(): + def __init__(self, enable_letter_box = False) -> None: + self.record_list = [] + self.enable_ltter_box = enable_letter_box + if self.enable_ltter_box is True: + self.letter_box_info_list = [] + else: + self.letter_box_info_list = None + + def letter_box(self, im, new_shape, pad_color=(0,0,0), info_need=False): + # Resize and pad image while meeting stride-multiple constraints + shape = im.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + # Scale ratio + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + + # Compute padding + ratio = r # width, height ratios + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding + + dw /= 2 # divide padding into 2 sides + dh /= 2 + + if shape[::-1] != new_unpad: # resize + im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=pad_color) # add border + + if self.enable_ltter_box is True: + self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, ratio, ratio, dw, dh, pad_color)) + if info_need is True: + return im, ratio, (dw, dh) + else: + return im + + def direct_resize(self, im, new_shape, info_need=False): + shape = im.shape[:2] + h_ratio = new_shape[0]/ shape[0] + w_ratio = new_shape[1]/ shape[1] + if self.enable_ltter_box is True: + self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, w_ratio, h_ratio, 0, 0, (0,0,0))) + im = cv2.resize(im, (new_shape[1], new_shape[0])) + return im + + def get_real_box(self, box, in_format='xyxy'): + bbox = copy(box) + if self.enable_ltter_box == True: + # unletter_box result + if in_format=='xyxy': + bbox[:,0] -= self.letter_box_info_list[-1].dw + bbox[:,0] /= self.letter_box_info_list[-1].w_ratio + bbox[:,0] = np.clip(bbox[:,0], 0, self.letter_box_info_list[-1].origin_shape[1]) + + bbox[:,1] -= self.letter_box_info_list[-1].dh + bbox[:,1] /= self.letter_box_info_list[-1].h_ratio + bbox[:,1] = np.clip(bbox[:,1], 0, self.letter_box_info_list[-1].origin_shape[0]) + + bbox[:,2] -= self.letter_box_info_list[-1].dw + bbox[:,2] /= self.letter_box_info_list[-1].w_ratio + bbox[:,2] = np.clip(bbox[:,2], 0, self.letter_box_info_list[-1].origin_shape[1]) + + bbox[:,3] -= self.letter_box_info_list[-1].dh + bbox[:,3] /= self.letter_box_info_list[-1].h_ratio + bbox[:,3] = np.clip(bbox[:,3], 0, self.letter_box_info_list[-1].origin_shape[0]) + return bbox + + def get_real_seg(self, seg): + #! fix side effect + dh = int(self.letter_box_info_list[-1].dh) + dw = int(self.letter_box_info_list[-1].dw) + origin_shape = self.letter_box_info_list[-1].origin_shape + new_shape = self.letter_box_info_list[-1].new_shape + if (dh == 0) and (dw == 0) and origin_shape == new_shape: + return seg + elif dh == 0 and dw != 0: + seg = seg[:, :, dw:-dw] # a[0:-0] = [] + elif dw == 0 and dh != 0 : + seg = seg[:, dh:-dh, :] + seg = np.where(seg, 1, 0).astype(np.uint8).transpose(1,2,0) + seg = cv2.resize(seg, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_LINEAR) + if len(seg.shape) < 3: + return seg[None,:,:] + else: + return seg.transpose(2,0,1) + + def add_single_record(self, image_id, category_id, bbox, score, in_format='xyxy', pred_masks = None): + if self.enable_ltter_box == True: + # unletter_box result + if in_format=='xyxy': + bbox[0] -= self.letter_box_info_list[-1].dw + bbox[0] /= self.letter_box_info_list[-1].w_ratio + + bbox[1] -= self.letter_box_info_list[-1].dh + bbox[1] /= self.letter_box_info_list[-1].h_ratio + + bbox[2] -= self.letter_box_info_list[-1].dw + bbox[2] /= self.letter_box_info_list[-1].w_ratio + + bbox[3] -= self.letter_box_info_list[-1].dh + bbox[3] /= self.letter_box_info_list[-1].h_ratio + # bbox = [value/self.letter_box_info_list[-1].ratio for value in bbox] + + if in_format=='xyxy': + # change xyxy to xywh + bbox[2] = bbox[2] - bbox[0] + bbox[3] = bbox[3] - bbox[1] + else: + assert False, "now only support xyxy format, please add code to support others format" + + def single_encode(x): + from pycocotools.mask import encode + rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0] + rle["counts"] = rle["counts"].decode("utf-8") + return rle + + if pred_masks is None: + self.record_list.append({"image_id": image_id, + "category_id": category_id, + "bbox":[round(x, 3) for x in bbox], + 'score': round(score, 5), + }) + else: + rles = single_encode(pred_masks) + self.record_list.append({"image_id": image_id, + "category_id": category_id, + "bbox":[round(x, 3) for x in bbox], + 'score': round(score, 5), + 'segmentation': rles, + }) + + def export_to_json(self, path): + with open(path, 'w') as f: + json.dump(self.record_list, f) + diff --git a/src/rk3576_yolov8_airborne/requirements.txt b/src/rk3576_yolov8_airborne/requirements.txt new file mode 100644 index 0000000..22afcc3 --- /dev/null +++ b/src/rk3576_yolov8_airborne/requirements.txt @@ -0,0 +1,5 @@ +numpy +opencv-python +fastapi +uvicorn +python-multipart \ No newline at end of file diff --git a/src/rk3576_yolov8_airborne/rknn-toolkit-lite2-packages/rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl b/src/rk3576_yolov8_airborne/rknn-toolkit-lite2-packages/rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl new file mode 100644 index 0000000..202237e --- /dev/null +++ b/src/rk3576_yolov8_airborne/rknn-toolkit-lite2-packages/rknn_toolkit_lite2-2.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bda74f1179e15fccb8726054a24898982522784b65bb340b20146955d254e800 +size 569160 diff --git a/src/rk3576_yolov8_airborne/video/test.mp4 b/src/rk3576_yolov8_airborne/video/test.mp4 new file mode 100644 index 0000000..bcaee49 --- /dev/null +++ b/src/rk3576_yolov8_airborne/video/test.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:61bea4a2bd769aaab6402cf9a4a5781454e8a9230f3c5a1fab35940ea04d62d8 +size 20033216 diff --git a/src/rk3576_yolov8_airborne/web_detection.py b/src/rk3576_yolov8_airborne/web_detection.py new file mode 100644 index 0000000..5e7b1e0 --- /dev/null +++ b/src/rk3576_yolov8_airborne/web_detection.py @@ -0,0 +1,864 @@ +import os +import sys +import cv2 +import argparse +import time +import numpy as np +import threading +from fastapi import FastAPI, Response, UploadFile, File, Form, HTTPException +from fastapi.responses import StreamingResponse, FileResponse +import uvicorn +import shutil +from typing import Optional, List + +# Import shared utilities +from py_utils.coco_utils import COCO_test_helper + +# Try to import RKNN-Toolkit-Lite2 +try: + from rknnlite.api import RKNNLite + RKNN_LITE_AVAILABLE = True +except ImportError: + RKNN_LITE_AVAILABLE = False + print("Warning: RKNN-Toolkit-Lite2 not available, using fallback") + +# Constants +OBJ_THRESH = 0.25 +NMS_THRESH = 0.45 +IMG_SIZE = (640, 640) # (width, height) + +# VisDrone classes (11 categories) +DEFAULT_CLASSES = ( + "pedestrian", "people", "bicycle", "car", "van", + "truck", "tricycle", "awning-tricycle", "bus", "motor", "others" +) + +CLASSES = DEFAULT_CLASSES + +def load_classes(path): + """Load classes from file, supports quoted comma-separated format.""" + global CLASSES + if not path or not os.path.exists(path): + CLASSES = DEFAULT_CLASSES + return + + try: + with open(path, 'r', encoding='utf-8') as f: + content = f.read().strip() + import re + items = re.findall(r'"([^"]*)"', content) + if items: + CLASSES = tuple(items) + print(f"Successfully loaded {len(CLASSES)} classes from {path}") + else: + items = [item.strip().strip('"') for item in content.split(',') if item.strip()] + if items: + CLASSES = tuple(items) + print(f"Loaded {len(CLASSES)} classes from {path} (fallback parsing)") + else: + print(f"Warning: No classes found in {path}, using default VisDrone classes") + CLASSES = DEFAULT_CLASSES + except Exception as e: + print(f"Error loading classes from {path}: {e}. Using default VisDrone classes") + CLASSES = DEFAULT_CLASSES + +# Dynamic config +class DetectionConfig: + def __init__(self): + self.obj_thresh = 0.25 + self.nms_thresh = 0.45 + self.lock = threading.Lock() + + def update(self, obj_thresh, nms_thresh): + with self.lock: + self.obj_thresh = obj_thresh + self.nms_thresh = nms_thresh + + def get(self): + with self.lock: + return self.obj_thresh, self.nms_thresh + +det_config = DetectionConfig() + +# --- Video analysis core --- +UPLOAD_DIR = "workspace/uploads" +OUTPUT_DIR = "workspace/outputs" +os.makedirs(UPLOAD_DIR, exist_ok=True) +os.makedirs(OUTPUT_DIR, exist_ok=True) + +class VideoAnalyzer: + def __init__(self, model=None, co_helper=None): + self.model = model + self.co_helper = co_helper + self.is_processing = False + self.progress = 0 + self.current_video = "" + self.error_msg = "" + self._stop_event = threading.Event() + self._thread = None + + def set_engine(self, model, co_helper): + self.model = model + self.co_helper = co_helper + + def start_analysis(self, input_path, output_path): + if self.is_processing: + return False + self._stop_event.clear() + self._thread = threading.Thread(target=self._process_video, args=(input_path, output_path)) + self._thread.daemon = True + self._thread.start() + return True + + def _process_video(self, input_path, output_path): + self.is_processing = True + self.progress = 0 + self.error_msg = "" + self.current_video = os.path.basename(input_path) + + cap = cv2.VideoCapture(input_path) + if not cap.isOpened(): + self.error_msg = f"Error: Cannot open video {input_path}" + self.is_processing = False + return + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if total_frames <= 0: + self.error_msg = "Error: Invalid total frames" + self.is_processing = False + cap.release() + return + + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) + + frame_idx = 0 + try: + while not self._stop_event.is_set(): + ret, frame = cap.read() + if not ret: + break + + if self.model and self.co_helper: + processed_img = preprocess_frame(frame, self.co_helper) + outputs = self.model.run(processed_img) + + if outputs is not None: + obj, nms = det_config.get() + boxes, classes, scores = post_process_with_thresh(outputs, obj, nms) + if boxes is not None: + draw(frame, self.co_helper.get_real_box(boxes), scores, classes) + + out.write(frame) + frame_idx += 1 + self.progress = int((frame_idx / total_frames) * 100) + + except Exception as e: + self.error_msg = f"Process error: {str(e)}" + finally: + cap.release() + out.release() + self.is_processing = False + if not self.error_msg: + self.progress = 100 + + def stop(self): + self._stop_event.set() + +video_analyzer = VideoAnalyzer() + +# --- FastAPI core --- +app = FastAPI(title="Airborne Object Detection (RK3576 NPU)") + +@app.get("/api/config") +async def get_config(): + obj, nms = det_config.get() + return {"obj_thresh": obj, "nms_thresh": nms} + +@app.post("/api/config") +async def update_config(config: dict): + det_config.update(config.get("obj_thresh", 0.25), config.get("nms_thresh", 0.45)) + return {"status": "success"} + +# --- Video analysis API --- +@app.post("/api/video/upload") +async def upload_video(file: UploadFile = File(...)): + file_path = os.path.join(UPLOAD_DIR, file.filename) + with open(file_path, "wb") as buffer: + shutil.copyfileobj(file.file, buffer) + return {"filename": file.filename, "status": "uploaded"} + +@app.get("/api/video/list") +async def list_videos(): + uploads = os.listdir(UPLOAD_DIR) if os.path.exists(UPLOAD_DIR) else [] + outputs = os.listdir(OUTPUT_DIR) if os.path.exists(OUTPUT_DIR) else [] + return {"uploads": uploads, "outputs": outputs} + +@app.post("/api/video/analyze") +async def analyze_video(filename: str = Form(...)): + input_path = os.path.join(UPLOAD_DIR, filename) + if not os.path.exists(input_path): + raise HTTPException(status_code=404, detail="Video not found") + + cap = cv2.VideoCapture(input_path) + if not cap.isOpened(): + raise HTTPException(status_code=400, detail="Cannot open video file") + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + cap.release() + + name_base = os.path.splitext(filename)[0] + output_filename = f"{name_base}_{width}x{height}_results.mp4" + output_path = os.path.join(OUTPUT_DIR, output_filename) + + success = video_analyzer.start_analysis(input_path, output_path) + if success: + return {"status": "started", "output": output_filename} + else: + return {"status": "error", "message": "Already processing another video"} + +@app.get("/api/video/status") +async def get_analysis_status(): + return { + "is_processing": video_analyzer.is_processing, + "progress": video_analyzer.progress, + "current_video": video_analyzer.current_video, + "error": video_analyzer.error_msg + } + +@app.get("/api/video/download/{filename}") +async def download_video(filename: str): + file_path = os.path.join(OUTPUT_DIR, filename) + if not os.path.exists(file_path): + raise HTTPException(status_code=404, detail="File not found") + return FileResponse(file_path, media_type='video/mp4', filename=filename) + +# Global variables for model access in API +_global_model = None +_global_co_helper = None + +@app.post("/api/models/yolov8/predict") +async def predict( + file: Optional[UploadFile] = File(None), + video: Optional[UploadFile] = File(None), + timestamp: Optional[float] = Form(None), + realtime: Optional[bool] = Form(False), + conf: Optional[float] = Form(None), + iou: Optional[float] = Form(None) +): + if _global_model is None or _global_co_helper is None: + return {"success": False, "message": "Model not initialized"} + + try: + img = None + source_info = "" + + # 1. Priority: uploaded image + if file: + contents = await file.read() + nparr = np.frombuffer(contents, np.uint8) + img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + source_info = "uploaded image" + + # 2. Priority: uploaded video with timestamp + elif video: + import tempfile + with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp: + tmp.write(await video.read()) + tmp_path = tmp.name + + cap = cv2.VideoCapture(tmp_path) + if cap.isOpened(): + if timestamp is not None: + cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000) + ret, frame = cap.read() + if ret: + img = frame + source_info = f"video frame at {timestamp if timestamp else 0}s" + cap.release() + os.unlink(tmp_path) + + # 3. Priority: realtime camera frame + if img is None: + img = frame_buffer.get_raw_frame() + source_info = "realtime camera frame" + + if img is None: + return {"success": False, "message": "No valid input source found (image, video, or camera)"} + + h, w = img.shape[:2] + + # Preprocess + input_img = preprocess_frame(img, _global_co_helper) + + # Inference + outputs = _global_model.run(input_img) + + # Use request params or global config + current_obj_thresh, current_nms_thresh = det_config.get() + target_conf = conf if conf is not None else current_obj_thresh + target_iou = iou if iou is not None else current_nms_thresh + + # Postprocess + boxes, classes, scores = post_process_with_thresh(outputs, target_conf, target_iou) + + predictions = [] + if boxes is not None and len(boxes) > 0: + real_boxes = _global_co_helper.get_real_box(boxes) + for box, score, cl in zip(real_boxes, scores, classes): + predictions.append({ + "class": CLASSES[cl], + "confidence": float(score), + "box": { + "x1": int(box[0]), + "y1": int(box[1]), + "x2": int(box[2]), + "y2": int(box[3]) + } + }) + + return { + "success": True, + "source": source_info, + "predictions": predictions, + "image": { + "width": w, + "height": h + } + } + except Exception as e: + return {"success": False, "message": str(e)} + +class FrameBuffer: + def __init__(self): + self.frame = None + self.raw_frame = None + self.lock = threading.Lock() + + def set_frame(self, frame, raw_frame=None): + with self.lock: + self.frame = frame + if raw_frame is not None: + self.raw_frame = raw_frame + + def get_frame(self): + with self.lock: + return self.frame + + def get_raw_frame(self): + with self.lock: + return self.raw_frame.copy() if self.raw_frame is not None else None + +frame_buffer = FrameBuffer() + +@app.get("/api/video_feed") +async def video_feed(): + def generate(): + while True: + frame = frame_buffer.get_frame() + if frame is not None: + yield (b'--frame\r\n' + b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n') + time.sleep(0.03) + return StreamingResponse(generate(), media_type="multipart/x-mixed-replace; boundary=frame") + +@app.get("/") +async def index(): + return Response(content=""" + + + Airborne Object Detection - RK3576 NPU + + + +
+

Airborne Object Detection (RK3576 NPU)

+ +
+
Real-time Detection
+
Local Video Analysis
+
+ +
+
+ +
+ +
+
+ +
+ + 0.25 +
+
+ +
+ +
+ + 0.45 +
+
+
+
+ +
+
+

Analyze Local Video

+
+ + + +
+ + + +
+ + + + + + + +
File NameAction
+
+
+
+ +

VisDrone YOLOv8s | Streaming via FastAPI + MJPEG | Port: 8000

+
+ + + + + """, media_type="text/html") + +def run_fastapi(host, port): + print("\n" + "="*50, flush=True) + print("Registered Routes:", flush=True) + for route in app.routes: + if hasattr(route, "methods"): + print(f"Path: {route.path:35} | Methods: {route.methods}", flush=True) + print("="*50 + "\n", flush=True) + sys.stdout.flush() + + uvicorn.run(app, host=host, port=port, log_level="info", log_config=None) + +# --- Inference logic (YOLOv8 DFL + NMS post-processing) --- + +def post_process_with_thresh(outputs, obj_thresh, nms_thresh): + """YOLOv8 post-processing with DFL decoding and NMS.""" + if outputs is None: + return None, None, None + + boxes, scores, classes_conf = [], [], [] + defualt_branch = 3 + pair_per_branch = len(outputs) // defualt_branch + for i in range(defualt_branch): + boxes.append(box_process(outputs[pair_per_branch * i])) + classes_conf.append(outputs[pair_per_branch * i + 1]) + scores.append(np.ones_like(outputs[pair_per_branch * i + 1][:, :1, :, :], dtype=np.float32)) + + def sp_flatten(_in): + ch = _in.shape[1] + _in = _in.transpose(0, 2, 3, 1) + return _in.reshape(-1, ch) + + boxes = np.concatenate([sp_flatten(_v) for _v in boxes]) + classes_conf = np.concatenate([sp_flatten(_v) for _v in classes_conf]) + scores = np.concatenate([sp_flatten(_v) for _v in scores]) + + # filter_boxes logic + scores_flat = scores.reshape(-1) + class_max_score = np.max(classes_conf, axis=-1) + classes = np.argmax(classes_conf, axis=-1) + _class_pos = np.where(class_max_score * scores_flat >= obj_thresh) + + scores = (class_max_score * scores_flat)[_class_pos] + boxes = boxes[_class_pos] + classes = classes[_class_pos] + + if len(classes) == 0: + return None, None, None + + nboxes, nclasses, nscores = [], [], [] + for c in set(classes): + inds = np.where(classes == c)[0] + b = boxes[inds] + s = scores[inds] + + # NMS + x = b[:, 0] + y = b[:, 1] + w = b[:, 2] - b[:, 0] + h = b[:, 3] - b[:, 1] + areas = w * h + order = s.argsort()[::-1] + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x[i], x[order[1:]]) + yy1 = np.maximum(y[i], y[order[1:]]) + xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) + yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) + w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) + h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) + inter = w1 * h1 + ovr = inter / (areas[i] + areas[order[1:]] - inter) + inds_nms = np.where(ovr <= nms_thresh)[0] + order = order[inds_nms + 1] + + if len(keep) > 0: + nboxes.append(b[keep]) + nclasses.append(np.full(len(keep), c)) + nscores.append(s[keep]) + + if not nboxes: + return None, None, None + + return np.concatenate(nboxes), np.concatenate(nclasses), np.concatenate(nscores) + +def dfl(position): + """Distribution Focal Loss decoding.""" + n, c, h, w = position.shape + p_num = 4 + mc = c // p_num + y = position.reshape(n, p_num, mc, h, w) + + y_max = np.max(y, axis=2, keepdims=True) + y_exp = np.exp(y - y_max) + y_softmax = y_exp / np.sum(y_exp, axis=2, keepdims=True) + + acc_matrix = np.arange(mc, dtype=np.float32).reshape(1, 1, mc, 1, 1) + y = np.sum(y_softmax * acc_matrix, axis=2) + return y + +def box_process(position): + """Decode box predictions from DFL output.""" + grid_h, grid_w = position.shape[2:4] + col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h)) + col = col.reshape(1, 1, grid_h, grid_w) + row = row.reshape(1, 1, grid_h, grid_w) + grid = np.concatenate((col, row), axis=1) + stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1) + position = dfl(position) + box_xy = grid +0.5 -position[:,0:2,:,:] + box_xy2 = grid +0.5 +position[:,2:4,:,:] + xyxy = np.concatenate((box_xy*stride, box_xy2*stride), axis=1) + return xyxy + +def post_process(input_data): + obj, nms = det_config.get() + return post_process_with_thresh(input_data, obj, nms) + +def draw(image, boxes, scores, classes): + for box, score, cl in zip(boxes, scores, classes): + top, left, right, bottom = [int(_b) for _b in box] + cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2) + cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score), + (top, left - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) + +class RKNNLiteModel: + def __init__(self, model_path): + if not RKNN_LITE_AVAILABLE: + raise ImportError("RKNN-Toolkit-Lite2 is not available") + if not os.path.exists(model_path): + raise FileNotFoundError(f"RKNN model file not found: {model_path}") + self.rknn_lite = RKNNLite() + print(f'Loading RKNN model from {model_path}...', flush=True) + sys.stdout.flush() + ret = self.rknn_lite.load_rknn(model_path) + if ret != 0: + raise Exception(f"Load RKNN model failed with error code: {ret}") + print('Initializing runtime...', flush=True) + sys.stdout.flush() + ret = self.rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2) + if ret != 0: + raise Exception(f"Init runtime failed with error code: {ret}") + print('RKNN model loaded successfully', flush=True) + sys.stdout.flush() + + def run(self, inputs): + try: + if len(inputs.shape) == 3: + inputs = np.expand_dims(inputs, axis=0) + if inputs.dtype != np.uint8: + inputs = inputs.astype(np.uint8) + return self.rknn_lite.inference(inputs=[inputs]) + except Exception as e: + print(f"Inference error: {e}") + return None + + def release(self): + if hasattr(self, 'rknn_lite'): + self.rknn_lite.release() + +def preprocess_frame(frame, co_helper): + img = co_helper.letter_box(im=frame.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0]), pad_color=(0,0,0)) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img + +def main(): + parser = argparse.ArgumentParser(description='Airborne Object Detection on RK3576 (Web Preview Mode)') + parser.add_argument('--model_path', type=str, required=True, help='RKNN model path') + parser.add_argument('--camera_id', type=int, default=1, help='Camera device ID (default: 1 for /dev/video1)') + parser.add_argument('--video_path', type=str, help='Path to video file (overrides camera_id)') + parser.add_argument('--class_path', type=str, help='Path to class_config.txt file for dynamic category loading') + parser.add_argument('--host', type=str, default='0.0.0.0', help='Web server host') + parser.add_argument('--port', type=int, default=8000, help='Web server port') + args = parser.parse_args() + + if not RKNN_LITE_AVAILABLE: + print("Error: RKNN-Toolkit-Lite2 is not available.") + return + + # Load custom classes + if args.class_path: + load_classes(args.class_path) + + global _global_model, _global_co_helper + # Initialize model + model = RKNNLiteModel(args.model_path) + co_helper = COCO_test_helper(enable_letter_box=True) + + _global_model = model + _global_co_helper = co_helper + video_analyzer.set_engine(model, co_helper) + + # Start web server thread + web_thread = threading.Thread(target=run_fastapi, args=(args.host, args.port), daemon=True) + web_thread.start() + print(f"Web Preview started at http://{args.host}:{args.port}", flush=True) + sys.stdout.flush() + + # If camera_id is -1, run in pure Web mode (video file analysis only) + if args.camera_id == -1: + print("Running in Video Analysis Mode. Access Web UI to process local videos.", flush=True) + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("Interrupted by user") + finally: + model.release() + return + + # Open video source + if args.video_path: + cap = cv2.VideoCapture(args.video_path) + else: + cap = cv2.VideoCapture(args.camera_id) + cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) + + if not cap.isOpened(): + print(f"Error: Cannot open video source (ID: {args.camera_id if not args.video_path else args.video_path})") + return + + fps_counter = 0 + try: + while True: + ret, frame = cap.read() + if not ret: + if args.video_path: + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + continue + break + + # Inference + processed_img = preprocess_frame(frame, co_helper) + start_time = time.time() + outputs = model.run(processed_img) + inference_time = time.time() - start_time + + if outputs is not None: + boxes, classes, scores = post_process(outputs) + if boxes is not None: + draw(frame, co_helper.get_real_box(boxes), scores, classes) + + # FPS + inf_fps = 1.0 / inference_time if inference_time > 0 else 0 + fps_counter = 0.9 * fps_counter + 0.1 * inf_fps if fps_counter > 0 else inf_fps + cv2.putText(frame, f'NPU FPS: {fps_counter:.1f}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) + + # Update web frame buffer + _, buffer = cv2.imencode('.jpg', frame) + frame_buffer.set_frame(buffer.tobytes(), frame) + + time.sleep(0.01) + + except KeyboardInterrupt: + print("Interrupted by user") + finally: + cap.release() + model.release() + +if __name__ == '__main__': + main() \ No newline at end of file