无人机数据集UAVDT转yolo数据集格式
将UAVDT数据集转换成yolo格式。
·
近期用到无人机UAVDT,需要在yolo上进行训练,下载后分为多目标和单目标两种,本次使用多目标数据进行训练,即对UAV-benchmark-MOTD_v1.0中的xx_gt_whole.txt真值文件进行修改。
查看了几篇修改过程:将UAVDT数据集格式转成VOC、YOLO格式(最终版)_uavdt数据集没标签-CSDN博客
都要先转换成COCO再到VOC最后再到yolo,不太适用于我,于是根据他们的教程进行修改,以下是我的转换过程:
1.处理数据集中的图像和相应的ground truth标签,将xx_gt_whole.txt按照图像文件的顺序保存到相应的目录中。生成目录如下:

def process_dataset(data_path):
data_num = 0
for name in os.listdir(data_path):
# print('name:', name)
data_num += 1
img_path = os.path.join(data_path, name)
gt_path = os.path.join(data_path.replace('UAV-benchmark-M', 'UAV-benchmark-MOTD_v1.0/GT'), name + '_gt_whole.txt')
save_gt_path = os.path.join(data_path, name, 'gt')
photo_count = 0
for file_name in os.listdir(img_path):
_, extension = os.path.splitext(file_name)
if extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
photo_count += 1
print("Total photos:", photo_count)
if not os.path.exists(save_gt_path):
os.makedirs(save_gt_path)
with open(gt_path, 'r') as file:
lines = file.readlines()
for i in range(1, photo_count + 1):
file_path1 = os.path.join(save_gt_path, 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)
# print('img_num={}\nend!!!'.format(data_num))
2.将所有图片和标注分别放在/img和/gt的文件夹中,并修改为一一对应的名字。

def img_and_gt(data_path, save_path):
os.makedirs(save_path, exist_ok=True)
# 目标图片文件夹路径
save_img_folder = os.path.join(save_path, 'img')
os.makedirs(save_img_folder, exist_ok=True)
# 目标标签文件夹路径
save_gt_folder = os.path.join(save_path, 'gt')
os.makedirs(save_gt_folder, exist_ok=True)
# 遍历指定文件夹中的所有文件夹
for folder_name in tqdm(os.listdir(data_path)):
# 构建当前文件夹中的图片文件夹路径
img_folder_path = os.path.join(data_path, folder_name)
# 用于计数的变量
img_count = 1
gt_count = 1
# 遍历图片文件夹中的所有文件
for filename in os.listdir(img_folder_path):
# 获取文件的扩展名
_, extension = os.path.splitext(filename)
if extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
# 获取文件的完整路径
file_path = os.path.join(img_folder_path, filename)
# 生成新的文件名
new_filename = f"{folder_name}_{img_count}{extension}"
# 构建目标文件的完整路径
destination_file = os.path.join(save_img_folder, new_filename)
# 移动文件到目标文件夹
# print(file_path,'/n',destination_file)
shutil.copy(file_path, destination_file)
# 更新计数器
# print('num:',img_count)
img_count += 1
# 构建当前文件夹中的标签文件夹路径
gt_folder_path = os.path.join(data_path, folder_name, "gt")
# 如果标签文件夹存在
if os.path.exists(gt_folder_path):
dir = os.listdir(gt_folder_path)
dir.sort(key=lambda x: int(x[:-4]))
# 遍历标签文件夹中的所有文件
for filename in dir:
# 获取文件的完整路径
file_path = os.path.join(gt_folder_path, filename)
# 生成新的文件名
new_filename = f"{folder_name}_{gt_count}.txt"
# 构建目标文件的完整路径
destination_file = os.path.join(save_gt_folder, new_filename)
# 移动文件到目标文件夹
shutil.copy(file_path, destination_file)
# 更新计数器
gt_count += 1
3.将所有图片分为训练集和测试集。
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 tqdm(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.copy(source_path, target_path)
4.为训练集分配对应的标签文件。
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 tqdm(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)
5.将标签文件格式转为yolo格式。
def convert_uavdt_to_yolo(img_folder, uavdt_labels_folder, yolo_labels_folder):
for uavdt_file in tqdm(os.listdir(uavdt_labels_folder)):
if uavdt_file.endswith('.txt'):
uavdt_labels_file = os.path.join(uavdt_labels_folder, uavdt_file)
os.makedirs(yolo_labels_folder, exist_ok=True)
yolo_labels_file = os.path.join(yolo_labels_folder, uavdt_file)
name, _ = os.path.splitext(uavdt_file)
img_file = os.path.join(img_folder, f"{name}.jpg")
with Image.open(img_file) as img:
img_width, img_height = img.size
with open(uavdt_labels_file, 'r') as f_in, open(yolo_labels_file, 'w') as f_out:
for line in f_in:
data = line.strip().split(',')
frame_index, target_id, bbox_left, bbox_top, bbox_width, bbox_height, out_of_view, occlusion, object_category = data
x_center = (float(bbox_left) + float(bbox_width) / 2) / img_width
y_center = (float(bbox_top) + float(bbox_height) / 2) / img_height
width = float(bbox_width) / img_width
height = float(bbox_height) / img_height
yolo_format = f"{object_category} {x_center} {y_center} {width} {height}"
f_out.write(yolo_format + '\n')
6.查看结果是否正确。
def visualize_yolo_labels(image_path, yolo_labels_path):
# 读取图像
image = cv2.imread(image_path)
print(image.shape)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 读取YOLO格式的标签文件
with open(yolo_labels_path, 'r') as f:
for line in f:
data = line.strip().split(' ')
object_category, x_center, y_center, width, height = data
# 计算边界框坐标
x_center, y_center, width, height = map(float, [x_center, y_center, width, height])
img_h, img_w, _ = image.shape
x_center *= img_w
y_center *= img_h
width *= img_w
height *= img_h
x1, y1, x2, y2 = int(x_center - width/2), int(y_center - height/2), int(x_center + width/2), int(y_center + height/2)
# 绘制边界框
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
cv2.putText(image, object_category, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

总共有三类:1:'car',2:'truck',3:'bus'
使用时:
# 调用函数
data_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M'
print('start step 1')
process_dataset(data_path)
save_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo'
print('start step 2')
img_and_gt(data_path, save_path)
img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\img'
train_img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_img'
test_img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\test_img'
print('start step 3')
split_dataset(img_dir, train_img_dir, test_img_dir)
txt_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\gt'
train_gt_uav_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_uav'
print('start step 4')
copy_txt_files(train_img_dir, txt_dir, train_gt_uav_dir)
train_gt_yolo_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_yolo'
print('start step 5')
convert_uavdt_to_yolo(train_img_dir, train_gt_uav_dir, train_gt_yolo_dir)
image_path = r'data/UAV-benchmark-M-yolo/train_img/M0201_181.jpg'
yoly_labels_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_yolo\M0201_181.txt'
visualize_yolo_labels(image_path, yoly_labels_path)
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)