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

news2025/1/31 23:50:25

前言

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

这将是我们本文的目标——在自定义数据集上训练不同的 YOLO NAS 模型。
在这里插入图片描述

YOLO-NAS训练

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

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

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

训练 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具有原始图像名称的目录中。
获得结果后，笔记本通过在预测图像上重叠地面实况注释来显示一组图像。
最后一步将告诉我们训练模型错过了哪些对象以及模型是否做出了错误的预测。

让我们通过可视化一些推理预测来开始我们的分析。
在这里插入图片描述