yolo-nas无人机高空红外热数据小目标检测(教程+代码)

前言

  • YOLO-NAS是目前最新的YOLO目标检测模型。
  • 从一开始,它就在准确性方面击败了所有其他 YOLO 模型。
  • 与之前的 YOLO 模型相比,预训练的 YOLO-NAS 模型能够以更高的准确度检测更多目标。
  • 但是我们如何在自定义数据集上训练 YOLO NAS?

这将是我们本文的目标------在自定义数据集上训练不同的 YOLO NAS 模型。

YOLO-NAS训练

YOLO-NAS 的主要主张是它可以比以前的模型更好地检测更小的物体。尽管我们可以运行多个推理实验来分析结果,但在具有挑战性的数据集上对其进行训练将使我们有更好的理解。为此,我们将使用三个可用的预训练 YOLO-NAS 模型运行四个训练实验。为此,我们选择无人机热成像检测数据集。

在实验过程中,我们将遍历 YOLO-NAS 的完整训练流程。

  1. 用于训练 YOLO NAS 的物体检测数据集
  2. 在自定义数据集上训练 YOLO NAS
  3. 微调 YOLO NAS 模型
  4. 使用经过训练的 YOLO NAS 模型对测试图像进行推理
  5. YOLO NAS 训练模型视频推理结果
  6. 结论

训练 YOLO NAS 的物体检测数据集

用于训练 YOLO NAS 的物体检测数据集

我们先来熟悉一下无人机高空红外热数据集。

它包含夜间无人机热图像。鉴于无人机的高空记录,大多数物体看起来都很小。这使得该数据集对于大多数目标检测模型来说都难以解决。然而,它是完美的自定义数据集来训练 YOLO-NAS 以检查其在小物体上的准确性。

该数据集包含 5 个对象类别的 2898 张热图像:

  • 自行车
  • 其他车辆
  • 不在乎

数据集已包含训练、验证和测试分割。有 2008 个训练样本、287 个验证样本和 571 个测试样本。该数据集已经以 YOLO 注释格式存在。

以下是数据集中的一些未注释的地面实况图像。

很明显,除了汽车之外,如果没有适当的注释,人眼无法看到地面上的其他物体。

要了解每个对象的位置,请查看一些带注释的图像

接下来,我们将深入研究本文的编码部分。下载本文的代码后,您将发现三个笔记本。

YOLO_NAS_Fine_Tuning.ipynb

YOLO_NAS_Large_Fine_Tuning.ipynb

inference.ipynb

YOLO_NAS_Fine_Tuning.ipynb我们将非常详细地浏览这些笔记本。这两个包含在自定义数据集上训练 YOLO NAS 以及稍后使用经过训练的模型运行推理所需的所有步骤。培训笔记本包含下载数据集的代码。

以下代码将训练三个 YOLO NAS 模型:

YOLO NAS (小)

YOLO NAS m(中型)

YOLO NAS l (大)

在开始之前,您可以安装super-gradients我们在整个训练和推理过程中需要的软件包。尽管笔记本包含执行此操作的命令,您也可以使用以下命令安装它:

复制代码
pip install

数据集下载和目录结构

接下来的几个代码块下载数据集并将其解压到当前目录,我们将在此处跳过。所有笔记本和数据集都存在于父数据集目录中,其结构如下

复制代码
hit-uav
├── dataset.yaml
├── images
│   ├── test
│   ├── train
│   └── val
└── labels
    ├── test
    ├── train
    └── val

YOLO NAS模型训练

由于我们正在训练三个不同的模型,因此我们需要稍微自动化该过程。我们可以定义一个包含三个模型名称的列表,并根据该列表设置检查点目录。这还将加载适当的模型,因为列表中的模型名称与 API 中的模型名称相匹配super-gradients。

复制代码
models_to_train = [
    'yolo_nas_s',
    'yolo_nas_m',
    'yolo_nas_l'
]
 
CHECKPOINT_DIR = 'checkpoints'
 
for model_to_train in models_to_train:
    trainer = Trainer(
        experiment_name=model_to_train, 
        ckpt_root_dir=CHECKPOINT_DIR
    )
 
    model = models.get(
        model_to_train, 
        num_classes=len(dataset_params['classes']), 
        pretrained_weights="coco"
    )
 
    trainer.train(
        model=model, 
        training_params=train_params, 
        train_loader=train_data, 
        valid_loader=val_data
    )

三个训练实验将依次运行,所有模型检查点将保存在各自的目录中。

YOLO NAS 训练参数

在我们开始微调过程之前,训练参数是最重要的组成部分。这是我们定义要训练的纪元数、要监控的验证指标以及学习率等的地方。、

复制代码
models_to_train = [
    'yolo_nas_s',
    'yolo_nas_m',
    'yolo_nas_l'
]
 
CHECKPOINT_DIR = 'checkpoints'
 
for model_to_train in models_to_train:
    trainer = Trainer(
        experiment_name=model_to_train, 
        ckpt_root_dir=CHECKPOINT_DIR
    )
 
    model = models.get(
        model_to_train, 
        num_classes=len(dataset_params['classes']), 
        pretrained_weights="coco"
    )
 
    trainer.train(
        model=model, 
        training_params=train_params, 
        train_loader=train_data, 
        valid_loader=val_data
    )

微调结果

YOLO NAS 模型对测试图像进​​行推理

该数据集包含一个测试分割,我们保留该测试分割用于推理目的。您可以执行笔记本中的代码单元inference.ipynb来运行推理实验。它促成了一些事情:

首先,它从检查点目录加载经过最佳训练的 YOLO NAS 权重。

然后它对测试图像运行推理。执行此操作时,代码会将推理结果保存在inference_results/images具有原始图像名称的目录中。

获得结果后,笔记本通过在预测图像上重叠地面实况注释来显示一组图像。

最后一步将告诉我们训练模型错过了哪些对象以及模型是否做出了错误的预测。

让我们通过可视化一些推理预测来开始我们的分析。

相关推荐
小雪狼12 小时前
RV1126 RKNN环境搭建记录
rnn·yolo
nju_spy14 小时前
计算机视觉 - 物体检测(二)单阶段:YOLO系列 + SSD
人工智能·yolo·目标检测·计算机视觉·ssd·r-cnn·端到端检测
码猩15 小时前
YOLO通用无人机目标检测框架
人工智能·yolo·目标检测
Hcoco_me15 小时前
YOLO入门教程(番外):计算机视觉数学、编程基础
人工智能·yolo·计算机视觉
berling001 天前
【论文阅读 | WACV 2025 | MCOR:通过跨模态信息互补和余弦相似性通道重采样模块增强的多光谱目标检测】
论文阅读·人工智能·目标检测
IT古董2 天前
【第五章:计算机视觉-项目实战之目标检测实战】2.目标检测实战:中国交通标志检测-(4)YOLOv8训练与测试
yolo·目标检测·计算机视觉
berling002 天前
【论文阅读 | TGRS 2025 | DHANet:用于多模态无人机目标检测的双流分层交互网络】
论文阅读·目标检测·无人机
pop_opo_2 天前
使用 Python + Pygame 键盘控制无人机(AirSim)
python·无人机·pygame
一朵小红花HH2 天前
SimpleBEV:改进的激光雷达-摄像头融合架构用于三维目标检测
论文阅读·人工智能·深度学习·目标检测·机器学习·计算机视觉·3d
云卓SKYDROID2 天前
无人机航电系统散热技术要点
人工智能·无人机·材质·高科技·云卓科技