YOLO通用无人机目标检测框架

YOLO通用目标检测框架

python 复制代码

import argparse
from pathlib import Path

from ultralytics import YOLO


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run YOLOv8 inference on images/video/camera")
    parser.add_argument(
        "--weights",
        type=str,
        default=str(Path("runs/bird/exp/weights/best.pt")),
        help="Path to trained weights (.pt)",
    )
    parser.add_argument(
        "--source",
        type=str,
        required=True,
        help="Input source: path to image/folder/video or camera index (e.g. 0)",
    )
    parser.add_argument("--imgsz", type=int, default=640, help="Inference image size")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    parser.add_argument("--iou", type=float, default=0.7, help="NMS IoU threshold")
    parser.add_argument("--device", type=str, default="cpu", help="Device: 'cpu' or '0' for GPU 0")
    parser.add_argument("--save", action="store_true", help="Save visualized predictions")
    parser.add_argument("--save_txt", action="store_true", help="Save predictions to .txt files")
    parser.add_argument("--save_crop", action="store_true", help="Save cropped detections")
    parser.add_argument("--show", action="store_true", help="Display a window with results")
    parser.add_argument("--project", type=str, default="runs/infer", help="Project dir for outputs")
    parser.add_argument("--name", type=str, default="predict", help="Run name for outputs")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    # Support numeric webcam indices
    source: str | int
    try:
        source = int(args.source)
    except ValueError:
        source = args.source

    model = YOLO(args.weights)
    model.predict(
        source=source,
        imgsz=args.imgsz,
        conf=args.conf,
        iou=args.iou,
        device=args.device,
        save=args.save,
        save_txt=args.save_txt,
        save_crop=args.save_crop,
        show=args.show,
        project=args.project,
        name=args.name,
        exist_ok=True,
        verbose=True,
    )


if __name__ == "__main__":
    main()

简介

YOLO通用目标检测框架是一个模块化、可扩展的目标检测解决方案，基于YOLO系列算法构建。该框架支持快速替换数据集并进行模型训练，提供从数据准备到模型部署的完整工作流程。

特性

🚀 多模型支持: 支持YOLOv5/YOLOv8等多个版本
📦 模块化设计: 清晰的架构，便于维护和扩展
🔄 完整流程: 涵盖数据准备、模型训练、推理和部署
⚡ 高性能: 支持GPU加速和模型优化
🌐 多种接口: 提供命令行、Python API和Web API
📊 可视化: 内置训练监控和结果可视化

目录结构

复制代码

yolosb/
├── models/              # 模型相关文件
├── datasets/            # 数据集管理
├── src/                 # 项目源代码
└── docs/                # 运行说明文档

快速开始

安装依赖

bash 复制代码

# 克隆项目
git clone <repository-url>
cd yolosb

# 创建虚拟环境
python -m venv venv
source venv/bin/activate  # Linux/macOS
# or
venv\Scripts\activate     # Windows

# 安装依赖
pip install -r requirements.txt

准备数据集

bash 复制代码

# 复制数据集模板
cp -r datasets/templates/yolo datasets/my_dataset

# 按照YOLO格式组织您的数据

模型训练

bash 复制代码

python src/scripts/train.py \
  --config models/configs/training/yolov8n.yaml \
  --data datasets/my_dataset/dataset.yaml \
  --epochs 100

目标检测

bash 复制代码

python src/scripts/detect.py \
  --model models/weights/trained/best.pt \
  --source input.jpg \
  --output results

项目源码地址：https://download.csdn.net/download/weixin_43050480/92043612