主要是将UAVDT数据集转换成COCO格式的目标检测数据集。
首先要将gt_whole.txt里面的DET数据按照片名称进行分割，保证分割后的txt文件中包含该照片的全部目标信息。相应的代码如下：
import os
#文件夹的路径
folder_path = r'E:\UAVDT_M\UAV-benchmark-M\UAV-benchmark-M\M1401'
folder_path1 = os.path.join(folder_path ,'gt')#存放分割后的txt文件的路径
folder_path2 = os.path.join(folder_path , 'img1')#照片的路径，主要用于统计有多少张照片，方便创建txt文件
folder_path3 = os.path.join(folder_path , 'gt', 'gt_whole.txt')#存放标签的txt文件
photo_count = 0
for file_name in os.listdir(folder_path2):
# 使用os.path.splitext()函数获取文件扩展名
_, extension = os.path.splitext(file_name)
# 如果文件是照片文件（例如.jpg或.png），则增加计数器
if extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
photo_count += 1
print("Total photos:", photo_count)
if not os.path.exists(folder_path):
os.makedirs(folder_path)
with open(folder_path3,'r') as file:
lines = file.readlines()
for i in range(1,photo_count+1):
file_path1 = os.path.join(folder_path1, str(i) + '.txt')
with open(file_path1, 'w') as target_file:
for line in lines:
data = line.split(',')
if data[0] == str(i):
target_file.write(line)
接下来是将各文件夹中的照片和txt文件按照依次递增的顺序存放到一个文件夹中，方便后面进行格式的转换，相应的代码如下：
import os
import shutil
folder_path = r'E:\UAVDT_M\UAV-benchmark-M\UAV-benchmark-M'
destination_folder = r'E:\UAVDT_M\UAV-benchmark-M\all\ann'
# 获取文件夹中的所有项目
items = os.listdir(folder_path)
count = 1
# 筛选出文件夹名称
folder_names = [item for item in items if os.path.isdir(os.path.join(folder_path, item))]
# 打印文件夹名称
for folder_name in folder_names:
folder_path2 = os.path.join(folder_path ,folder_name, 'gt')
print(folder_path2)
for filename in os.listdir(folder_path2):
# 获取文件的完整路径
folder_path3 = os.path.join(folder_path2, filename)
# 获取文件的扩展名
_, extension = os.path.splitext(filename)
# 生成新的文件名
new_filename = f'image_{count}{extension}'
# 构建目标文件的完整路径
destination_file = os.path.join(destination_folder, new_filename)
# 复制文件到目标文件夹
shutil.copy(folder_path3, destination_file)
# 更新计数器
count += 1
将全部的文件按照一定的比例进行切分，这儿设置的比例是2：8，可以直接修改参数split_ratio，相应的代码如下：
import os
import random
import shutil
def split_dataset(source_dir, train_dir, test_dir, split_ratio=0.8):
# 创建训练集和测试集文件夹
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# 获取所有图片文件的列表
image_files = [f for f in os.listdir(source_dir) if f.endswith('.jpg') or f.endswith('.png')]
# 随机化文件列表
random.shuffle(image_files)
# 计算训练集和测试集的数量
num_train = int(len(image_files) * split_ratio)
num_test = len(image_files) - num_train
# 将文件移动到相应的文件夹中
for i, image_file in enumerate(image_files):
source_path = os.path.join(source_dir, image_file)
if i < num_train:
target_path = os.path.join(train_dir, image_file)
else:
target_path = os.path.join(test_dir, image_file)
shutil.move(source_path, target_path)
# 设置源文件夹和目标文件夹路径
source_directory = r'E:\UAVDT_M\UAV-benchmark-M\all\img'
train_directory = r'E:\UAVDT_M\UAV-benchmark-M\all\train'
test_directory = r'E:\UAVDT_M\UAV-benchmark-M\all\test'
# 划分数据集并移动文件
split_dataset(source_directory, train_directory, test_directory)
分割完成后，按照分割后的图片名称，将相对应的标签文件进行分割，相应的代码如下：
import os
import shutil
def copy_txt_files(image_folder, txt_folder, output_folder):
# 创建输出文件夹
os.makedirs(output_folder, exist_ok=True)
# 获取图片文件夹中所有图片文件的名称列表
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg') or f.endswith('.png')]
# 遍历图片文件夹中的图片文件
for image_file in image_files:
# 构建对应的txt文件路径
txt_file = os.path.splitext(image_file)[0] + '.txt'
txt_path = os.path.join(txt_folder, txt_file)
# 检查对应的txt文件是否存在
if os.path.exists(txt_path):
# 复制txt文件到输出文件夹中
shutil.copy(txt_path, output_folder)
# 设置图片文件夹路径、txt文件夹路径和输出文件夹路径
image_folder = r'E:\UAVDT_M\UAV-benchmark-M\all\test'
txt_folder = r'E:\UAVDT_M\UAV-benchmark-M\all\ann'
output_folder = r'E:\UAVDT_M\UAV-benchmark-M\all\test_ann'
# 将txt文件复制到新的文件夹中
copy_txt_files(image_folder, txt_folder, output_folder)
最后将分割完成的标签文件由txt转换成为json进行保存，相应的代码如下：
import json
import os
# 创建一个空的COCO格式数据结构
coco_data = {
"info": {},
"licenses": [],
"categories": [],
"images": [],
"annotations": []
}
# 读取txt文件并解析数据
def parse_txt_file(txt_file):
with open(txt_file, 'r') as f:
lines = f.readlines()
annotations = []
for line in lines:
# 解析每一行数据，格式为 <frame_index>,<target_id>,<bbox_left>,<bbox_top>,<bbox_width>,<bbox_height>,<out-of-view>,<occlusion>,<object_category>
frame_index, target_id, bbox_left, bbox_top, bbox_width, bbox_height, out_of_view, occlusion, object_category = map(str.strip, line.split(','))
annotations.append({
"bbox": [int(bbox_left), int(bbox_top), int(bbox_width), int(bbox_height)],
"category_id": int(object_category)
})
return annotations
# 遍历txt文件夹
txt_folder = r'E:\UAVDT_M\UAV-benchmark-M\all\test_ann'
for filename in os.listdir(txt_folder):
if filename.endswith('.txt'):
txt_file = os.path.join(txt_folder, filename)
# 解析txt文件
annotations = parse_txt_file(txt_file)
# 将图像信息添加到COCO数据结构中
image_info = {
"id": len(coco_data["images"]) + 1,
"file_name": os.path.splitext(filename)[0] + '.jpg' # 假设txt文件与图像文件同名
}
coco_data["images"].append(image_info)
# 将目标标注信息添加到COCO数据结构中
for annotation in annotations:
annotation_info = {
"id": len(coco_data["annotations"]) + 1,
"image_id": image_info["id"],
"bbox": annotation["bbox"],
"category_id": annotation["category_id"],
"iscrowd": 0 # 假设目标不是一个crowd
}
coco_data["annotations"].append(annotation_info)
# 保存COCO格式数据为JSON文件
with open(r'E:\UAVDT_M\UAV-benchmark-M\all\test.json', 'w') as f:
json.dump(coco_data, f)
