# tensorRT端侧部署

**Repository Path**: altria1122/tensor-rt-end-side-deployment

## Basic Information

- **Project Name**: tensorRT端侧部署
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-01-28
- **Last Updated**: 2026-01-29

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# TensorRT Edge-Side Deployment

This project is an end-to-end solution integrating model training, edge inference, and web management. It enables users to train YOLO models using Python and deploy them on a high-performance edge inference engine based on TensorRT and C++, while unified management and monitoring are provided through a Node.js service.

## Features

1. **Powerful Model Training Framework (`python-train/`)**
   * Supports training of **YOLOv8** object detection models.
   * Supports **Knowledge Distillation**, including feature distillation and logit distillation, for model compression or incremental learning (class expansion).
   * Supports **COCO** and **custom YOLO** format datasets.
   * Includes rich data augmentation strategies (Mosaic, MixUp, HSV, random affine transformations, etc.).

2. **High-Performance Inference Engine (`cpp-engine/`)**
   * C++ inference backend based on **TensorRT**, supporting FP16/INT8 acceleration.
   * Supports dynamic input dimensions (Dynamic Shapes) and batch inference.
   * Integrates **Mediamtx** (formerly rtsp-simple-server) for receiving and outputting RTSP streams.
   * Supports concurrent inference from multiple camera streams.

3. **Web Management Service (`src/`)**
   * RESTful API service built on **Node.js (Express)**.
   * **WebSocket** real-time communication for pushing detection results and alert notifications.
   * Provides comprehensive interfaces for **model management**, **camera management**, **data management**, and **training control**.
   * Integrated Swagger API documentation.

## Project Structure

```text
tensor-rt-end-side-deployment/
├── config/                  # Configuration files directory
├── cpp-engine/              # C++ inference engine core code
│   ├── include/             # Header files
│   ├── src/
│   │   ├── command_handler.cpp # Command handling logic (ZMQ)
│   │   ├── inference/       # TensorRT engine wrapper
│   │   ├── inference_worker.cpp # Inference worker thread and video stream processing
│   │   └── zmq_server.cpp   # ZMQ communication server
│   ├── mediamtx/            # Embedded RTSP server configuration
│   ├── push_streams.sh      # Stream pushing script
│   └── CMakeLists.txt       # Build configuration
├── python-train/            # Python training framework
│   ├── configs/             # Training configuration files (YAML)
│   ├── data/                # Dataset processing and augmentation
│   ├── losses/              # Loss functions (YOLO Loss & Distill Loss)
│   ├── models/              # YOLOv8 model definitions (Backbone, Neck, Head)
│   ├── trainer/             # Trainer logic
│   └── train.py             # Training entry point
├── src/                     # Node.js backend service
│   ├── routes/              # API routes (camera, inference, model, train)
│   ├── services/            # Core service modules
│   │   ├── inferenceEngine.js # Communication with C++ engine
│   │   ├── streamManager.js   # Video stream management
│   │   └── trainManager.js    # Training process management
│   └── app.js               # Service entry point
└── logs/                    # Log files directory
```

## Quick Start

### 1. Python Training Environment (Optional)

If you need to train or fine-tune models, set up the Python environment.

**Install dependencies:**
```bash
cd python-train
pip install -r requirements.txt
```

**Start training:**
```bash
# Standard training
python train.py --config configs/train_default.yaml

# Knowledge distillation training
python distill.py --config configs/train_default.yaml
```
*Refer to `python-train/README.md` or configuration files in the `configs/` directory for detailed configuration instructions.*

### 2. C++ Inference Engine Build

The inference engine depends on CUDA, TensorRT, and ZMQ.

1. Install dependencies (Ubuntu example):
    ```bash
    sudo apt-get install cmake libzmq3-dev
    # Install CUDA and TensorRT (download and install manually)
    ```
2. Build the project:
    ```bash
    cd cpp-engine
    mkdir build && cd build
    cmake ..
    make -j4
    ```

### 3. Node.js Service Deployment

The web management platform is the core of the system, responsible for orchestrating the C++ engine and managing data.

**Start the service:**
```bash
cd src
npm install
node app.js
```

After startup, the service runs on port `3000` by default. Access `http://localhost:3000/api-docs` to view the Swagger API documentation.

## Core API Overview

*   **Model Management (`/api/model`)**: Upload model weights (`.pt`) and automatically or manually convert them to ONNX and TensorRT Engine formats.
*   **Inference Control (`/api/inference`)**: Initialize the engine, start/stop inference, and set batch size.
*   **Camera Management (`/api/camera`)**: Add/remove RTSP camera streams and test connection status.
*   **Data Management (`/api/data`)**: View inference result images, delete data, and export datasets for training.

## Architecture Design

1. **Data Flow**:  
   Camera RTSP stream → **StreamManager** → **C++ Inference Worker** → **TensorRT Engine** → **Detection Results** → **WebSocket** → **Frontend Display**

2. **Control Flow**:  
   Frontend request → **Node.js API** → **ZMQ Request** → **C++ CommandHandler** → **Inference Worker**

## Contributing

Issues and pull requests are welcome.

## License

This project is licensed under an open-source license. See the LICENSE file in the project root for details.