# smallYolo

**Repository Path**: hejun5448/smallYolo

## Basic Information

- **Project Name**: smallYolo
- **Description**: yolo功能实现，用作于中间平台使用。
- **Primary Language**: Java
- **License**: Apache-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 1
- **Created**: 2026-09-24
- **Last Updated**: 2026-09-24

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# smallYolo

> 将 YOLO 核心推理能力封装为 Spring Boot Web 服务，提供开箱即用的 Web UI 与开放 API。  
> 支持图片/视频检测、在线标注训练、YOLO 模型训练导出 ONNX、API Key 鉴权等能力。

[![Gitee](https://img.shields.io/badge/Gitee-smallYolo-C71D23?logo=gitee&logoColor=white)](https://gitee.com/channnlang/smallYolo.git)

---

## ✨ 特性

- **图片目标检测** — 上传图片返回检测框、标签、置信度及标注后图片（Base64）
- **视频异步检测** — 上传视频后后台逐帧推理，前端轮询进度，输出 H.264 标注视频在线播放/下载
- **在线标注训练** — Web 页面直接上传图片、Canvas 画框标注、填写类别，即可生成数据集或直接训练
- **YOLO 模型训练** — 上传 YOLO 格式数据集 zip，后端编排本地 Python（ultralytics）异步训练，自动导出 ONNX 并注册到模型列表
- **模型热管理** — `models/` 目录扫描，按需懒加载 `OrtSession` 并缓存；输入尺寸自动读取 ONNX 形状
- **API Key 鉴权** — 开放接口（`/api/open/*`）支持 API Key 管理，明文 Key 仅创建时返回一次，后续只存 SHA-256 哈希
- **多输出结构自动识别** — 无需关心模型类型：

  | 模型 | 输出结构 |
  |---|---|
  | YOLOv5 | `[1, 25200, 5+C]` |
  | YOLOv7 (end2end) | `[n, 7]` |
  | YOLOv8 / v9 / v11 | `[1, 4+C, 8400]` |

---

## 🛠 技术栈

- Java 17 / Spring Boot 3.3.5
- ONNX Runtime Java API
- OpenCV（图像预处理与视频编解码）
- 单文件 HTML/CSS/JS 前端
- 文件持久化（无数据库依赖）

---

## 🚀 快速开始

### 1. 环境要求

- JDK 17+
- Maven 3.9+
- （可选训练功能）Python 3.9+ 及可联网环境

### 2. 启动服务

```bash
cd smallYolo
export JAVA_HOME=$(/usr/libexec/java_home -v 17)
mvn spring-boot:run
```

浏览器打开：http://localhost:8080

### 3. 启用训练功能（可选）

```bash
./scripts/setup-python-env.sh
```

该脚本会自动创建 `python-env/` 虚拟环境并安装 `ultralytics`，约占用 2GB 磁盘空间。

> **注意**：代码目录不要包含中文路径。使用 GPU 需安装 CUDA/cuDNN，并将 `pom.xml` 中 `onnxruntime` 替换为 `onnxruntime_gpu`，同时在 `ModelManager#createSession` 中开启 `options.addCUDA(0)`。

---

## 📁 目录说明

```
smallYolo/
├── models/                 # 放置 .onnx 模型文件，启动后自动扫描加载
├── pretrained/             # 基础 .pt 模型（训练用），如 yolov8n.pt
├── training/               # 训练任务运行时工作目录（数据集解压、训练日志、中间产物）
├── python-env/             # Python 虚拟环境（训练依赖，可删除后重新 setup）
├── api-keys/               # API Key 持久化文件（keys.json，只存哈希）
├── scripts/                # setup-python-env.sh、train_runner.py 等脚本
├── src/main/java/          # Spring Boot 后端
└── src/main/resources/static/index.html  # 单文件前端
```

---

## ⚙️ 配置

`src/main/resources/application.yml`：

```yaml
server:
  port: 8080

yolo:
  model-dir: models
  default-model: yolov8s.onnx
  conf-threshold: 0.35
  nms-threshold: 0.55
  api-key-dir: api-keys
  model-labels:
    smoke: [no_smoke, smoke]
```

---

## 🔌 API 概览

### 开放接口（需 API Key，CORS 已开放）

| 接口 | 方法 | 说明 |
|---|---|---|
| `/api/open/detect` | POST | 图片检测，返回纯 JSON 数据 |
| `/api/open/detect/base64` | POST | Base64 图片检测 |
| `/api/keys/exists` | GET | 查询是否已创建 API Key |
| `/api/keys` | POST/GET/DELETE | 创建/列出/删除 API Key |

### Web UI 接口

| 接口 | 说明 |
|---|---|
| `GET /api/models` | 模型列表、加载状态、默认阈值 |
| `POST /api/detect` | 图片检测，返回含标注图的结果 |
| `POST /api/video/detect` | 提交视频检测任务 |
| `GET /api/video/tasks/{id}` | 查询视频任务进度 |
| `POST /api/train/submit` | 提交模型训练任务 |
| `GET /api/train/tasks/{id}` | 查询训练任务进度 |
| `POST /api/annotation/export` | Canvas 标注导出 YOLO 数据集 zip |
| `POST /api/annotation/train` | Canvas 标注后直接训练 |

详细参数与 curl 示例请启动服务后点击页面右上角 **📖 API 文档** 查阅。

---

## 🧠 使用自己的模型

1. 将 `.onnx` 文件放入 `models/` 目录（或修改 `yolo.model-dir`）
2. 自定义标签：在 `application.yml` 的 `yolo.model-labels` 中按文件名包含匹配配置
3. 模型输出需为上述三种结构之一

---

## 📌 注意事项

- 首次创建 API Key 时无需鉴权（bootstrap 模式），创建第一个 Key 后开放接口将强制校验 `X-API-Key`
- 训练产物（数据集、runs）默认保留 6 小时后自动清理，导出的 `.onnx` 模型永久保留在 `models/`
- 视频任务结果保留 2 小时后自动清理
- CPU 推理适合验证与小规模使用；生产环境建议 GPU 并开启多线程/流水线优化

---

## 🔗 相关链接

- **Gitee 仓库**：https://gitee.com/channnlang/smallYolo.git

---

## 📄 License

MIT