AI-TOD 数据集整理

本文档用于归档 AI-TOD 遥感小目标检测数据集的完整整理流程，涵盖数据下载、数据处理、格式转换三个核心阶段，流程规范可复现，适配后续 YOLO/MMDetection 等框架的模型训练需求。

一、 数据下载（阶段 1）

1.1 下载所需数据集

AI-TOD 数据集的生成依赖两个核心数据源，需分别下载并归档：

     2.  AI-TOD_wo_xview 数据集（基础背景画布库）

• 下载地址：https://drive.google.com/drive/folders/1uNY_rcOO5LrWibXRY6l2dvqSbK6xikJp
• 归档内容：下载得到的压缩包、解压后的标注文件（.json 格式，存放至 annotations 目录）、图像文件（.png 格式，按 train/val/test 三级目录存放）
• 归档路径：/media/lfq/EAGET/AI-TOD/AI-TOD/aitod/（保持官方指定目录结构不变）
• 关键说明：该数据集是 AI-TOD 的基础背景库，图像数量分别为 train 7510 张、val 2804 张、test 9351 张，无需修改其目录结构，避免影响后续合成流程。

1.2 下载配套工具

1.3 环境准备

• 激活专属 Conda 环境：conda activate aitod（Python 3.7，适配 AI-TOD 相关工具依赖）
• 预装基础依赖：pip install numpy pillow tqdm shapely fiona（为后续数据处理、格式转换铺路）

二、 数据处理（阶段 2）

2.1 数据解压与目录规整

2.2 运行 AI-TOD 数据合成脚本

核心目标：将 xView 小目标素材嵌入 AI-TOD_wo_xview 背景图像，生成完整的 AI-TOD 数据集（9306 张，train/3704 张、val/935 张、test/4667 张）。

切换至脚本目录：cd /media/lfq/EAGET/AI-TOD/AI-TOD/

        1. 准备wwtool工具

git clone https://github.com/jwwangchn/wwtool.git
cd wwtool
python setup.py develop

        安装其他重要的包

cd ..
cd aitodtoolkit
pip install -r requirements.txt

运行运行合成脚本：

python generate_aitod_imgs.py

  2. 监控运行日志：无报错中断、无缺失文件提示，等待合成完成（耗时根据硬件配置有所差异）。    3. 验证合成结果：            查看输出目录 /media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets/，确认存在 train/val/test 三级目录。           随机打开若干图像，检查小目标嵌入自然，无明显拼接痕迹，图像无损坏。           核对图像数量：train 3704 张、val 935 张、test 4667 张，总数 9306 张，符合官方标准。

2.3 原始标注归档

合成完成后，将生成的原始标注（VOC 风格 .txt 格式，分布在 train/val/test 三级目录下）分别归档至对应目录的 ./voc_anno/ 文件夹，路径示例：

• train 集标注：./xview/xview_aitod_sets/train/voc_anno/
• val 集标注：./xview/xview_aitod_sets/val/voc_anno/
• test 集标注：./xview/xview_aitod_sets/test/voc_anno/确保标注文件与图像文件一一对应（同名，仅后缀不同），无缺失、无冗余。

三、 数据格式转换（阶段 3）

核心目标：将 VOC 风格 .txt 标注（xmin ymin xmax ymax class_name）转换为 YOLO 标准格式 .txt 标注（class_id x_center y_center w h，归一化至 0-1），适配 YOLOv5/YOLOv8 等框架训练。

3.1 步骤 1：提取全量类别（多目录批量提取）

import os

# ———————- 配置区（修改为你的三个标注目录路径）———————-
# 填写 train、val、test 三个目录的完整路径（按你的实际路径修改）
VOC_TXT_DIRS = [
"/media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets/trainval/labels",
"/media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets/test/labels",

]
# ——————————————————————————–

# 1. 初始化数据结构（去重+统计）
class_set = set() # 存储所有不重复类别
class_count_total = {} # 统计所有目录的总标注数量
class_count_per_dir = {} # 统计每个目录的标注数量（方便查看分布）

# 2. 遍历所有指定目录
for voc_txt_dir in VOC_TXT_DIRS:
# 初始化当前目录的类别统计
dir_name = os.path.basename(os.path.dirname(voc_txt_dir)) # 提取目录名（train/val/test）
class_count_per_dir[dir_name] = {}

# 检查目录是否存在
if not os.path.exists(voc_txt_dir):
print(f"警告：目录不存在，跳过 → {voc_txt_dir}")
continue

# 遍历当前目录下的所有 .txt 标注文件
voc_txt_files = [f for f in os.listdir(voc_txt_dir) if f.endswith(".txt")]
if not voc_txt_files:
print(f"警告：目录下无 .txt 标注文件，跳过 → {voc_txt_dir}")
continue

# 解析每个标注文件
for txt_file in voc_txt_files:
txt_path = os.path.join(voc_txt_dir, txt_file)
with open(txt_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if not line:
continue # 跳过空行

# 解析每行标注（格式：xmin ymin xmax ymax class_name）
parts = line.split()
if len(parts) != 5:
continue # 跳过格式错误的行

# 提取类别名称
class_name = parts[-1]

# 3. 加入全局集合（自动去重）
class_set.add(class_name)

# 4. 统计总数量
if class_name in class_count_total:
class_count_total[class_name] += 1
else:
class_count_total[class_name] = 1

# 5. 统计当前目录的数量
if class_name in class_count_per_dir[dir_name]:
class_count_per_dir[dir_name][class_name] += 1
else:
class_count_per_dir[dir_name][class_name] = 1

# 3. 转换为有序列表（方便分配 YOLO 类别 ID）
class_list = sorted(list(class_set))

# 4. 输出结果
print("=" * 60)
print(f"成功从 {len([d for d in VOC_TXT_DIRS if os.path.exists(d)])} 个有效目录中提取到 {len(class_list)} 个类别：")
print("=" * 60)

# 输出所有类别（带全局序号和总数量）
for idx, class_name in enumerate(class_list):
total_count = class_count_total.get(class_name, 0)
print(f"类别 ID {idx} → {class_name}（总标注数量：{total_count}）")

# 输出每个目录的类别分布（可选，方便查看数据均衡性）
print("\\n" + "=" * 60)
print("各目录类别标注数量分布：")
print("=" * 60)
for dir_name, class_count in class_count_per_dir.items():
print(f"\\n【{dir_name} 目录】")
for class_name in class_list:
dir_count = class_count.get(class_name, 0)
print(f" {class_name}：{dir_count} 条标注")

# 输出可直接复制的 CLASS_MAP
print("\\n" + "=" * 60)
print("直接可复制到 VOC → YOLO 转换脚本的 CLASS_MAP：")
print("=" * 60)
class_map_code = "CLASS_MAP = {\\n"
for idx, class_name in enumerate(class_list):
class_map_code += f' "{class_name}": {idx},\\n'
class_map_code = class_map_code.rstrip(",\\n") + "\\n}"
print(class_map_code)

3.2 步骤 2：批量转换标注格式

import os
from tqdm import tqdm

# ———————- 配置区（必须根据你的实际情况修改）———————-
# 1. 你的 VOC 风格 .txt 标注文件夹路径
input_voc_txt_dir = "/media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets/trainval/labels" # 替换为你的 VOC .txt 文件夹路径
# 2. YOLO 格式 .txt 标注输出文件夹路径
output_yolo_txt_dir = "/media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets/trainval/yolo-labels"
# 3. AI-TOD 图像尺寸（默认 1024×1024，无需修改，若不同请调整）
IMG_WIDTH, IMG_HEIGHT = 1024, 1024
# 4. 类别映射表（key：你的 VOC 标注中的文本类别，value：YOLO 数字 ID（从 0 开始））
# 示例：如果你的标注只有 "vehicle" 类，就只保留这一行；有多个类就依次添加
CLASS_MAP = {
"airplane": 0,
"ship": 1,
"storage-tank": 2,
"vehicle": 3
}

# ——————————————————————————–

# 1. 确保输出目录存在
os.makedirs(output_yolo_txt_dir, exist_ok=True)

# 2. 遍历所有 VOC 风格 .txt 标注文件
voc_txt_files = [f for f in os.listdir(input_voc_txt_dir) if f.endswith(".txt")]
if not voc_txt_files:
print("错误：未在输入目录中找到 .txt 标注文件")
exit(1)

for txt_file in tqdm(voc_txt_files, desc="转换 VOC → YOLO"):
# 2.1 拼接输入/输出文件路径
input_txt_path = os.path.join(input_voc_txt_dir, txt_file)
output_txt_path = os.path.join(output_yolo_txt_dir, txt_file)

# 2.2 读取 VOC 标注，转换并写入 YOLO 标注
with open(input_txt_path, 'r', encoding='utf-8') as f_in, \\
open(output_txt_path, 'w', encoding='utf-8') as f_out:

for line in f_in:
line = line.strip()
if not line:
continue # 跳过空行

# 解析 VOC 格式一行数据（xmin ymin xmax ymax class_name）
# 注意：如果你的数据有多余空格/制表符，split() 会自动处理
parts = line.split()
if len(parts) != 5:
print(f"警告：跳过无效行（格式错误）→ {line}")
continue

xmin, ymin, xmax, ymax, class_name = parts

# 2.3 转换为数值类型（像素坐标）
try:
xmin = float(xmin)
ymin = float(ymin)
xmax = float(xmax)
ymax = float(ymax)
except ValueError:
print(f"警告：跳过无效行（坐标非数字）→ {line}")
continue

# 2.4 检查类别是否在映射表中
if class_name not in CLASS_MAP:
print(f"警告：跳过未知类别 → {class_name}")
continue
class_id = CLASS_MAP[class_name]

# 2.5 转换为 YOLO 格式（归一化 + 中心坐标+宽高）
# 计算中心坐标
x_center = (xmin + xmax) / 2.0
y_center = (ymin + ymax) / 2.0
# 计算宽高
bbox_w = xmax – xmin
bbox_h = ymax – ymin
# 归一化（除以图像宽高，保留 6 位小数）
x_center_norm = round(x_center / IMG_WIDTH, 6)
y_center_norm = round(y_center / IMG_HEIGHT, 6)
bbox_w_norm = round(bbox_w / IMG_WIDTH, 6)
bbox_h_norm = round(bbox_h / IMG_HEIGHT, 6)

# 2.6 写入 YOLO 标注文件
f_out.write(f"{class_id} {x_center_norm} {y_center_norm} {bbox_w_norm} {bbox_h_norm}\\n")

# 3. 转换完成提示
print(f"\\n转换成功！YOLO 标注文件已保存至：{output_yolo_txt_dir}")
print(f"共处理 {len(voc_txt_files)} 个标注文件")

3.3 步骤 3：验证转换结果

3.4 步骤 4：生成 YOLO 框架配置文件

新建 aitod.yaml 配置文件，填写数据集路径、类别信息，归档至 ./xview/xview_aitod_sets/ 目录，内容示例：

path: /media/lfq/EAGET/AI-TOD/AI-TOD/aitodtoolkit/xview/xview_aitod_sets # 数据集根路径
train: train # 训练集图像目录（相对根路径）
val: val # 验证集图像目录（相对根路径）
test: test # 测试集图像目录（相对根路径）
nc: 4 # 类别总数（AI-TOD 官方 28 类）
names: ["airplane", "ship", "storage-tank","vehicle"] # 完整类别名称列表（与 CLASS_MAP 一一对应）

四、 最终归档与总结

总结