CMFADet复现(Linux)
·
前期准备
DroneVehicle数据集下载:
Train (code:ngar)
Validation (code:jnj6)
Test (code:tqwc)
数据集格式转换:
#!/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 &
更多推荐

所有评论(0)