目标检测-RT-DETR学习笔记（1）

RT-DETR官方代码给出了基于两种深度学习架构的代码，分别是pytorch版本和paddle版本，二者都是深度学习架构,paddle是百度团队开发的一种深度学习架构。

我们将介绍pytorch版本的。

第一步：环境配置

官方使用的是torch2.0.1版本，建议先在pytorch官网：download.pytorch.org/whl/cu118中下载gpu版本，然后离线安装torch，安装完成torch之后，在使用pip install -r requirements.txt安装其他依赖环境，安装完成requirements.txt中的模块后，还需要手动安装transformers ：pip install transformers。如果激活环境后，直接在虚拟环境中使用pip install -r requirements.txt命令安装的会是cpu版本的torch。

第二步：数据集转换

RT-DETR官方代码使用的是coco格式的数据集，如果你的数据集是yolo的txt格式的数据集，需要转换成coco格式的数据集。转换方法如下。

yolo数据集存放格式如下

包括一个images文件夹和一个labels文件夹，images文件夹下包括训练数据集和验证数据集，labels文件夹下包含训练数据集标签和验证数据集标签。

使用如下代码将yolo数据集转换成coco数据集标注格式。

import os
import json
from PIL import Image

# 设置数据集路径
dataset_path = "E:\\pdx\\PythonProject\\yolo2coco\\data"
images_path = "E:\\pdx\\PythonProject\\yolo2coco\\data\\images"
labels_path = "E:\\pdx\\PythonProject\\yolo2coco\\data\\labels"

# 类别映射
categories = [
    {"id": 0, "name": "HLB"},
    {"id": 1, "name": "KYB"},
    # 添加更多类别
]


# YOLO格式转COCO格式的函数
def convert_yolo_to_coco(x_center, y_center, width, height, img_width, img_height):
    x_min = (x_center - width / 2) * img_width
    y_min = (y_center - height / 2) * img_height
    width = width * img_width
    height = height * img_height
    return [x_min, y_min, width, height]


# 初始化COCO数据结构
def init_coco_format():
    return {
        "images": [],
        "annotations": [],
        "categories": categories
    }


# 处理每个数据集分区
for split in ['train','test', 'val']:
    coco_format = init_coco_format()
    annotation_id = 1

    for img_name in os.listdir(os.path.join(images_path, split)):
        if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
            img_path = os.path.join(images_path, split, img_name)
            label_path = os.path.join(labels_path, split, img_name.replace("JPG", "txt"))

            img = Image.open(img_path)
            img_width, img_height = img.size
            image_info = {
                "file_name": img_name,
                "id": len(coco_format["images"]) + 1,
                "width": img_width,
                "height": img_height
            }
            coco_format["images"].append(image_info)

            if os.path.exists(label_path):
                with open(label_path, "r") as file:
                    for line in file:
                        category_id, x_center, y_center, width, height = map(float, line.split())
                        bbox = convert_yolo_to_coco(x_center, y_center, width, height, img_width, img_height)
                        annotation = {
                            "id": annotation_id,
                            "image_id": image_info["id"],
                            "category_id": int(category_id) + 1,
                            "bbox": bbox,
                            "area": bbox[2] * bbox[3],
                            "iscrowd": 0
                        }
                        coco_format["annotations"].append(annotation)
                        annotation_id += 1

    # 为每个分区保存JSON文件
    with open(f"E:\\pdx\\PythonProject\\yolo2coco\\data\\coco\\{split}_coco_format.json", "w") as json_file:
        json.dump(coco_format, json_file, indent=4)

转换完成后coco数据集存放格式如下：

annotations文件夹中存放的是训练数据集和验证数据集的.json标注文件。

train和val中存放的是训练数据和测试数据图像。

第三步 进行训练

第一，修改rtdetr_pytorch/configs/dataset/coco_detection.yml中img_folder和ann_file为自己数据的地址。注意！！train_dataloader下的img_folder和ann_file修改为训练数据集的地址，val_dataloader下的img_folder和ann_file修改为验证数据集的地址。

第二，修改tools/train.py中config中default参数，将图片左侧的.yml文件地址写在default中

default="E:\\pdx\\PythonProject\\rtdetr_pytorch\\configs\\rtdetr\\rtdetr_r18vd_6x_coco.yml"