前期准备

DroneVehicle数据集下载:

Train (code:ngar)

Validation (code:jnj6)

Test (code:tqwc)

模型下载:Yooyoo95/CMFADet: A Cross-Modality Feature Adaptive Interaction Approach for RGB-Infrared Object Detection in Aerial Imagery

数据集格式转换:

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Convert the raw DroneVehicle dataset (VOC-style XML annotations with rotated
4-point polygons) into the layout required by CMFADet (multimodal RGB+IR
YOLO-OBB training):

    <out>/
        images/train|val|test/*.jpg     RGB images  (folder MUST be `images`)
        image/train|val|test/*.jpg      IR  images  (folder MUST be `image`)
        labels/train|val|test/*.txt     YOLO-OBB labels
        data_DV_CMFADet.yaml            dataset config

The folder names ``images`` / ``image`` are hard requirements of the CMFADet
loader:
  * the IR path is derived by replacing ``images`` with ``image``
  * the label path is derived by replacing ``images`` with ``labels``

Label format (one object per line, normalized to image width/height):
    OBB:  <class> x1 y1 x2 y2 x3 y3 x4 y4
    HBB:  <class> cx cy w h

Optional white-border cropping: the raw DroneVehicle images are 840x712 with
a 100 px white border on every side (content 640x512). Use ``--crop 100`` to
remove the border from the images and shift the labels accordingly.

Official DroneVehicle has 5 classes: car, truck, bus, van, freight_car.
A few annotations contain typos (``feright car``, ``feright_car``,
``feright``, ``truvk``) which are mapped to the correct class; the single
invalid ``*`` label is dropped.

Usage:
    python convert_dronevehicle_to_cmfadet.py \
        --src E:/temporary_e/DroneVehicle \
        --out E:/temporary_e/DroneVehicle_CMFADet \
        --crop 100
"""

import argparse
import shutil
import sys
import xml.etree.ElementTree as ET
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


CLASS_NAMES = ["car", "truck", "bus", "van", "freight_car"]

# Canonical class name -> class id (official DroneVehicle naming).
CLASS_MAP = {
    "car": 0,
    "truck": 1,
    "bus": 2,
    "van": 3,
    "freight_car": 4,
    # typos that appear in the official DroneVehicle annotations
    "feright_car": 4,
    "feright car": 4,
    "feright": 4,
    "truvk": 1,
}
# Any other name (e.g. the single "*" label) is dropped and reported.

# split -> (rgb folder, ir folder, xml folder)
SPLIT_DIRS = {
    "train": ("trainimg", "trainimgr", "trainlabel"),
    "val": ("valimg", "valimgr", "vallabel"),
    "test": ("testimg", "testimgr", "testlabel"),
}

POINT_TAGS = ["x1", "y1", "x2", "y2", "x3", "y3", "x4", "y4"]
EPS = 1e-6


def read_image_size(path: Path):
    """Read (width, height) of an image without loading the full array."""
    try:
        from PIL import Image

        with Image.open(path) as im:
            return im.size
    except Exception:
        return None


def polygon_area(points):
    """Shoelace area of a polygon (positive for clockwise, negative for CCW)."""
    area = 0.0
    n = len(points)
    for i in range(n):
        x1, y1 = points[i]
        x2, y2 = points[(i + 1) % n]
        area += x1 * y2 - x2 * y1
    return area / 2.0


def xml_to_label_line(xml_path: Path, rgb_img_path: Path, out_format: str,
                      crop=(0, 0, 0, 0)):
    """
    Convert one XML file to a list of label lines.

    Returns (lines, stats) where stats is a dict with per-file counters.
    Coordinates are normalized with the (cropped) image width/height; on any
    out-of-range result we fall back to the real image size.
    """
    stats = {"processed": 0, "dropped_unknown_cls": 0, "dropped_invalid": 0}
    lines = []
    left, right, top, bottom = crop

    tree = ET.parse(xml_path)
    root = tree.getroot()
    size_el = root.find("size")
    w = float(size_el.find("width").text) if size_el is not None else None
    h = float(size_el.find("height").text) if size_el is not None else None

    for obj in root.iter("object"):
        stats["processed"] += 1
        name_el = obj.find("name")
        if name_el is None or name_el.text is None:
            stats["dropped_unknown_cls"] += 1
            continue
        name = name_el.text.strip().lower()
        cls_id = CLASS_MAP.get(name)
        if cls_id is None:
            stats["dropped_unknown_cls"] += 1
            continue

        poly = obj.find("polygon")
        coords = None
        if poly is not None:
            coords = []
            ok = True
            for tag in POINT_TAGS:
                el = poly.find(tag)
                if el is None or el.text is None:
                    ok = False
                    break
                try:
                    coords.append(float(el.text))
                except ValueError:
                    ok = False
                    break
            if not ok or len(coords) != 8:
                coords = None
        else:
            # Some official annotations use an axis-aligned <bndbox> instead
            # of a rotated <polygon>; keep them as axis-aligned 4 corners.
            box = obj.find("bndbox")
            if box is not None:
                coords = []
                corners = [("xmin", "ymin"), ("xmax", "ymin"),
                           ("xmax", "ymax"), ("xmin", "ymax")]
                ok = True
                for x_tag, y_tag in corners:
                    x_el, y_el = box.find(x_tag), box.find(y_tag)
                    if x_el is None or y_el is None or x_el.text is None or y_el.text is None:
                        ok = False
                        break
                    try:
                        coords.append(float(x_el.text))
                        coords.append(float(y_el.text))
                    except ValueError:
                        ok = False
                        break
                if not ok or len(coords) != 8:
                    coords = None
        if coords is None:
            stats["dropped_invalid"] += 1
            continue

        pts = [(coords[i], coords[i + 1]) for i in range(0, 8, 2)]
        if abs(polygon_area(pts)) <= 0:
            stats["dropped_invalid"] += 1
            continue

        if w is None or h is None or h <= 0 or w <= 0:
            size = read_image_size(rgb_img_path)
            if size is None:
                stats["dropped_invalid"] += 1
                continue
            w, h = float(size[0]), float(size[1])

        cw = w - left - right
        ch = h - top - bottom
        if cw <= 0 or ch <= 0:
            stats["dropped_invalid"] += 1
            continue
        norm = [
            (coords[i] - (left if i % 2 == 0 else top)) / (cw if i % 2 == 0 else ch)
            for i in range(8)
        ]
        if min(norm) < -EPS or max(norm) > 1 + EPS:
            # Either the XML size disagrees with the real image or the object
            # is cut off by the border crop; renormalize with the actual
            # image size and then clip into the valid [0, 1] range.
            size = read_image_size(rgb_img_path)
            if size is not None:
                rw, rh = float(size[0]), float(size[1])
                cw = rw - left - right
                ch = rh - top - bottom
                if cw > 0 and ch > 0:
                    norm = [
                        (coords[i] - (left if i % 2 == 0 else top)) /
                        (cw if i % 2 == 0 else ch)
                        for i in range(8)
                    ]
        norm = [min(max(v, 0.0), 1.0) for v in norm]
        if all(v == 0.0 for v in norm):
            # The whole box ended up inside the cropped-away border.
            stats["dropped_invalid"] += 1
            continue

        if out_format == "obb":
            nums = [str(cls_id)] + [f"{v:.6g}" for v in norm]
        else:  # hbb
            xs = norm[0::2]
            ys = norm[1::2]
            cx = (min(xs) + max(xs)) / 2.0
            cy = (min(ys) + max(ys)) / 2.0
            bw = max(xs) - min(xs)
            bh = max(ys) - min(ys)
            nums = [str(cls_id), f"{cx:.6g}", f"{cy:.6g}", f"{bw:.6g}", f"{bh:.6g}"]
        lines.append(" ".join(nums))

    return lines, stats


def copy_image(src: Path, dst: Path, crop=(0, 0, 0, 0)):
    """Copy one image, optionally cropping the white border."""
    left, right, top, bottom = crop
    if left or right or top or bottom:
        from PIL import Image

        with Image.open(src) as im:
            w, h = im.size
            im.crop((left, top, w - right, h - bottom)).save(dst, quality=95)
        return None
    # resume-friendly: skip if identical size exists
    if dst.exists() and dst.stat().st_size == src.stat().st_size:
        return None
    shutil.copy2(src, dst)
    return None


def convert_split(src_root: Path, out_root: Path, split: str, out_format: str,
                  copy_images: bool, workers: int, crop=(0, 0, 0, 0)):
    rgb_dir, ir_dir, xml_dir = SPLIT_DIRS[split]
    src_rgb = src_root / split / rgb_dir
    src_ir = src_root / split / ir_dir
    src_xml = src_root / split / xml_dir

    dst_rgb = out_root / "images" / split
    dst_ir = out_root / "image" / split
    dst_lbl = out_root / "labels" / split
    dst_rgb.mkdir(parents=True, exist_ok=True)
    dst_ir.mkdir(parents=True, exist_ok=True)
    dst_lbl.mkdir(parents=True, exist_ok=True)

    xml_files = sorted(src_xml.glob("*.xml"))
    class_counter = Counter()
    dropped = Counter()
    label_count = 0
    empty_label_count = 0

    for i, xml_path in enumerate(xml_files, 1):
        stem = xml_path.stem
        rgb_src = src_rgb / f"{stem}.jpg"
        lines, stats = xml_to_label_line(xml_path, rgb_src, out_format, crop)
        if not lines:
            empty_label_count += 1
        label_count += len(lines)
        dropped["dropped_unknown_cls"] += stats["dropped_unknown_cls"]
        dropped["dropped_invalid"] += stats["dropped_invalid"]
        class_counter.update(int(line.split()[0]) for line in lines)
        (dst_lbl / f"{stem}.txt").write_text("\n".join(lines), encoding="utf-8")
        if i % 2000 == 0:
            print(f"    [{split}] converted {i}/{len(xml_files)} XMLs", flush=True)

    if copy_images:
        jobs = []
        with ThreadPoolExecutor(max_workers=workers) as pool:
            for src, dst in ((src_rgb, dst_rgb), (src_ir, dst_ir)):
                for p in src.glob("*.jpg"):
                    jobs.append(pool.submit(copy_image, p, dst / p.name, crop))
            done = 0
            for _ in as_completed(jobs):
                done += 1
                if done % 4000 == 0:
                    print(f"    [{split}] copied {done}/{len(jobs)} images", flush=True)

    return {
        "split": split,
        "images": len(xml_files),
        "labels": label_count,
        "empty": empty_label_count,
        "classes": dict(class_counter),
        "dropped": dict(dropped),
    }


def write_yaml(out_root: Path):
    cfg = f"""# DroneVehicle (RGB + Infrared) dataset for CMFADet
# Converted from the raw VOC/XML DroneVehicle annotations (100 px white
# border cropped, images 640x512).
path: {out_root.as_posix()}

## RGB images
train: images/train
val: images/val
test: images/test

## Infrared images
train_ir: image/train
val_ir: image/val
test_ir: image/test

## Classes (official DroneVehicle naming)
names:
  0: car
  1: truck
  2: bus
  3: van
  4: freight_car
"""
    (out_root / "data_DV_CMFADet.yaml").write_text(cfg, encoding="utf-8")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src", required=True, help="Raw DroneVehicle root, e.g. E:/temporary_e/DroneVehicle")
    parser.add_argument("--out", required=True, help="Output dataset root, e.g. E:/temporary_e/DroneVehicle_CMFADet")
    parser.add_argument("--format", choices=["obb", "hbb"], default="obb",
                        help="Label format: obb (4 corners, default) or hbb (xywh)")
    parser.add_argument("--skip-images", action="store_true",
                        help="Only convert labels + write the yaml, do not copy images")
    parser.add_argument("--workers", type=int, default=8, help="Threads used for image copies")
    parser.add_argument("--crop", type=int, default=0,
                        help="Crop N px white border on all four sides (e.g. 100)")
    parser.add_argument("--crop-left", type=int, default=0, help="Left crop px")
    parser.add_argument("--crop-right", type=int, default=0, help="Right crop px")
    parser.add_argument("--crop-top", type=int, default=0, help="Top crop px")
    parser.add_argument("--crop-bottom", type=int, default=0, help="Bottom crop px")
    args = parser.parse_args()

    crop = (args.crop_left or args.crop, args.crop_right or args.crop,
            args.crop_top or args.crop, args.crop_bottom or args.crop)

    src_root = Path(args.src)
    out_root = Path(args.out)
    if not src_root.is_dir():
        sys.exit(f"Source directory not found: {src_root}")
    if out_root.exists() and not out_root.is_dir():
        sys.exit(f"Output path exists and is not a directory: {out_root}")
    out_root.mkdir(parents=True, exist_ok=True)

    print(f"Converting DroneVehicle: {src_root} -> {out_root} "
          f"(format={args.format}, crop={crop})")
    reports = []
    for split in ("train", "val", "test"):
        print(f"  [{split}] starting ...", flush=True)
        r = convert_split(src_root, out_root, split, args.format,
                          copy_images=not args.skip_images, workers=args.workers,
                          crop=crop)
        reports.append(r)
        print(f"  [{split}] done: {r['images']} images, {r['labels']} labels, "
              f"{r['empty']} empty, dropped={r['dropped']}", flush=True)

    write_yaml(out_root)

    report_path = out_root / "conversion_report.txt"
    with report_path.open("w", encoding="utf-8") as f:
        f.write(f"src: {src_root}\n")
        f.write(f"out: {out_root}\n")
        f.write(f"label format: {args.format}\n\n")
        for r in reports:
            f.write(f"split {r['split']}: images={r['images']} "
                    f"objects={r['labels']} empty_label_files={r['empty']}\n")
            f.write(f"  class distribution: {r['classes']}\n")
            f.write(f"  dropped objects: {r['dropped']}\n")
    print(f"YAML + report written to {out_root}")


if __name__ == "__main__":
    main()

数据集目录结构为:

DroneVehicle/
├── image/          # 红外图像
│   ├── train/
│   ├── val/
│   └── test/
├── images/         # RGB 图像
│   ├── train/
│   ├── val/
│   └── test/
└── labels/         # 旋转框标签(8点归一化格式)
    ├── train/
    ├── val/
    └── test/

环境配置

conda create -n cmfadet python=3.10 -y  #创建虚拟环境
conda activate cmfadet  #激活环境
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121  #安装PyTorch (根据 CUDA 版本选择)
git clone https://github.com/Yooyoo95/CMFADet.git  #克隆代码仓库
cd CMFADet  #进入项目文件
#安装依赖
# 通用依赖
pip install opencv-python pandas matplotlib tqdm pyyaml scipy pillow requests psutil py-cpuinfo prettytable

# 必需的特殊依赖
pip install dill timm efficientnet_pytorch einops tensorboard thop

# 安装可变形卷积库 mmcv(务必匹配 PyTorch 和 CUDA 版本)
# 以 CUDA 12.1 + PyTorch 2.1 为例:
pip install mmcv==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.1/index.html

# 降级 NumPy 到 1.26.4(避免与 PyTorch 2.1 冲突)
pip install numpy==1.26.4 --force-reinstall

修改配置文件

1.修改 dataset/data_DV_Multimodel.yaml,将路径改为你的实际绝对路径,并正确指定 RGB 和红外目录

train: /你的路径/DroneVehicle/images/train   # RGB
val: /你的路径/DroneVehicle/images/val
test: /你的路径/DroneVehicle/images/test

train_ir: /你的路径/DroneVehicle/image/train   # 红外
val_ir: /你的路径/DroneVehicle/image/val
test_ir: /你的路径/DroneVehicle/image/test

names:
  0: car
  1: truck
  2: bus
  3: van
  4: freight_car

2.修改训练脚本 train.py。

将模型配置文件路径改为仓库内实际存在的多模态模型 yaml

model = YOLO('/你的路径/CMFADet/ultralytics/cfg/models/multimodal/Multi-SFEM-IRAFAB-CIF-obb-ATAH.yaml')

将数据配置文件路径改为上一步修改好的 yaml

data='/你的路径/CMFADet/dataset/data_DV_Multimodel.yaml'

确保超参数与论文一致

imgsz=640,
epochs=300,          # 注意拼写是 epochs,不是 pochs
batch=8,
optimizer='SGD',
amp=False,
project='runs/DroneVehicle/train',
name='CMFADet-obb-640',

启动训练

python train.py
#或者后台运行
nohup python train.py > train.log 2>&1 &

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐