【目标检测】d-fine模型部署

官网介绍显示，d-fine模型效果很好，例如下图中，非常模糊的人也能被识别出来。官网教程有些细节没写，这里补充一下。

1.数据格式

数据长这样（图中的unrrelated_pic无用）。

具体的格式可以由大模型生成。我这里有数据转化的代码，我的原始数据长这样。

在sample_detect文件夹下，请只关注images和labels两个文件夹，annotations是转化代码生成的文件。images文件夹下有train、test、val三个文件夹，直接放图片即可，labels文件夹下有对应的所有图片的txt标注文件，类似于yolo。

代码运行时，请修改CLASS_NAMES 、ROOT_DIR 和convert_txt_to_coco即可。

import os
import json
import cv2  # 使用 OpenCV 读取图片尺寸
from tqdm import tqdm  # 显示进度条，提供更好的用户体验

# --- 1. 配置信息 ---

# 你的类别名称，ID会从1开始分配（COCO标准中，0通常是背景）
# 请务必按照你的 a.txt 文件中的 class_id 顺序来填写
CLASS_NAMES = ["xxx"]

# 数据集根目录
ROOT_DIR = r'F:\yyy'


# --- 2. 核心转换函数 ---

def convert_txt_to_coco(root_dir, dataset_split):
    """
    将 YOLO 格式的 TXT 标注文件转换为 COCO JSON 格式。

    Args:
        root_dir (str): 数据集根目录，包含 images 和 labels 文件夹。
        dataset_split (str): 'train' 或 'val'，指定要转换的数据子集。
    """
    print(f"--- 开始转换 {dataset_split} 数据集 ---")

    # 定义输入路径
    img_dir = os.path.join(root_dir, 'images', dataset_split)
    txt_dir = os.path.join(root_dir, 'labels', dataset_split)

    # 定义输出文件路径
    output_dir = os.path.join(root_dir, 'annotations')
    os.makedirs(output_dir, exist_ok=True)
    json_path = os.path.join(output_dir, f'{dataset_split}.json')

    # --- 初始化 COCO JSON 结构 ---
    coco_data = {
        "info": {
            "description": f"Custom dataset converted from TXT to COCO format - {dataset_split}",
            "version": "1.0",
            "year": 2025,
            "contributor": "YourName",  # 请替换为你的名字
            "date_created": "2025/07/24"
        },
        "licenses": [],
        "images": [],
        "annotations": [],
        "categories": []
    }

    # ---填充类别信息 (categories) ---
    # COCO 的 category_id 通常从 1 开始
    for i, name in enumerate(CLASS_NAMES):
        coco_data['categories'].append({
            "id": i + 1,
            "name": name,
            "supercategory": "object"
        })

    # --- 核心转换逻辑 ---
    image_id_counter = 1
    annotation_id_counter = 1

    # 遍历所有图片文件
    image_files = [f for f in os.listdir(img_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

    for image_filename in tqdm(image_files, desc=f"处理 {dataset_split} 图片"):
        # 构造对应的 txt 文件路径
        base_filename = os.path.splitext(image_filename)[0]
        txt_filename = base_filename + '.txt'
        txt_filepath = os.path.join(txt_dir, txt_filename)

        # 读取图片获取尺寸
        img_path = os.path.join(img_dir, image_filename)
        try:
            # 使用 imread 读取图片，即使是灰度图也能正常获取宽高
            img = cv2.imread(img_path)
            if img is None:
                print(f"警告：无法读取图片 {img_path}，跳过该文件。")
                continue
            img_height, img_width, _ = img.shape
        except Exception as e:
            print(f"错误：读取图片 {img_path} 失败，原因为: {e}。跳过该文件。")
            continue

        # --- 填充图片信息 (images) ---
        coco_data['images'].append({
            "id": image_id_counter,
            "file_name": image_filename,
            "width": img_width,
            "height": img_height
        })

        # --- 填充标注信息 (annotations) ---
        if os.path.exists(txt_filepath):
            with open(txt_filepath, 'r') as f:
                for line in f.readlines():
                    try:
                        class_id, x_center_norm, y_center_norm, width_norm, height_norm = map(float,
                                                                                              line.strip().split())

                        # 将 YOLO 格式 (归一化) 转换为 COCO 格式 (绝对像素值)
                        # bbox: [x_min, y_min, width, height]
                        bbox_width = width_norm * img_width
                        bbox_height = height_norm * img_height
                        x_min = (x_center_norm * img_width) - (bbox_width / 2)
                        y_min = (y_center_norm * img_height) - (bbox_height / 2)

                        # 计算面积
                        area = bbox_width * bbox_height

                        coco_data['annotations'].append({
                            "id": annotation_id_counter,
                            "image_id": image_id_counter,
                            "category_id": int(class_id) + 1,  # TXT中class_id从0开始，COCO从1开始
                            "bbox": [round(x_min, 2), round(y_min, 2), round(bbox_width, 2), round(bbox_height, 2)],
                            "area": round(area, 2),
                            "iscrowd": 0,
                            "segmentation": []  # 对于bbox，segmentation通常为空
                        })
                        annotation_id_counter += 1

                    except ValueError:
                        print(f"警告：文件 {txt_filename} 中的行格式不正确: '{line.strip()}'")

        image_id_counter += 1

    # --- 保存为 JSON 文件 ---
    with open(json_path, 'w') as f:
        json.dump(coco_data, f, indent=4)

    print(f"--- 转换完成！JSON 文件已保存至: {json_path} ---")
    print(f"总计图片数: {len(coco_data['images'])}")
    print(f"总计标注数: {len(coco_data['annotations'])}")


# --- 3. 执行转换 ---
if __name__ == '__main__':
    # 转换训练集
    convert_txt_to_coco(ROOT_DIR, 'train')

    # # 转换验证集
    # convert_txt_to_coco(ROOT_DIR, 'val')

2.配置文件

你需要修改图中的文件，即“D-FINE-master/configs/dfine/dfine_hgnetv2_n_coco.yml”。你也可以选择其他尺寸的模型。

你需要修改代码中的第2、4行才能完成数据加载，他们都在configs文件夹下。batch_size需要在这几个文件中都去修改一下。num_classes这个参数也很坑，默认777，改了之后会报错。其他训练参数也在对应的文件中。

3.运行代码

完成配置后，即可执行下面的命令单卡部署

CUDA_VISIBLE_DEVICES=0,1 python  train.py  -c /home/xxx/代码/D-FINE-master/configs/dfine/dfine_hgnetv2_n_coco.yml --use-amp --seed=0 -t /home/xxx/代码/D-FINE-master/pretrain_model/dfine_n_coco.pth

双卡部署时，请参照原作者readme文件中，torchrun相关方法，并使用sudo nvidia-smi -pm 1确保多张显卡开启。

运行成功截图如下：

4.注意事项

推理的时候，在D-FINE-master/tools/inference文件夹下，需要配一个新环境，但是这个环境会与训练时候的环境冲突，请格外小心，做好备份或者另开一个环境。