最近需要将UAVDT数据集转VOC标注格式,想找个现成的代码直接用,结果发现咋都还要会员呢,算了自己写一个吧!
代码贴在这里可以直接用。
注意,使用这个代码前我将UAVDT图片文件全部加上了视频名的前缀(也就是图片的父文件夹的名字)这样可以避免图片重名,并把所有图片移动到了同一个文件夹中,如果不是这样那么代码的逻辑需要进行一定的修改。
import os
import xml.etree.ElementTree as ET
import cv2
def process_files(gt_path, images_path, output_path, class_names):
"""
处理 UAVDT-MOT 数据集,将 GT 文件转换为 VOC 格式的 XML 文件。
:param gt_path: GT 文件所在目录
:param images_path: 图像文件所在目录
:param output_path: 输出的 VOC 格式 XML 文件目录
"""
for file_name in os.listdir(gt_path):
if not file_name.endswith('gt_whole.txt'):
continue
video_name = file_name[:-13]
gt_file_path = os.path.join(gt_path, file_name)
for line in open(gt_file_path, 'r'):
datas = line.strip().split(',')
frame_id = int(datas[0])
xmin = int(float(datas[2]))
ymin = int(float(datas[3]))
xmax = xmin + int(float(datas[4]))
ymax = ymin + int(float(datas[5]))
class_id = int(datas[8])
calss_name = class_names[class_id - 1]
image_name = f"{video_name}_img{frame_id:06d}.jpg" # 我这里是因为我把UAVDT的图片名加上了视频名的前缀避免重复
image_path = os.path.join(images_path, image_name)
xml_path = os.path.join(output_path, f"{video_name}_img{frame_id:06d}.xml")
if not os.path.exists(xml_path):
# 创建 XML 根节点
annotation = ET.Element("annotation")
ET.SubElement(annotation, "folder").text = "VOC2007"
ET.SubElement(annotation, "filename").text = image_name
ET.SubElement(annotation, "path").text = image_path
# 添加图像尺寸信息
size = ET.SubElement(annotation, "size")
img_width, img_height = cv2.imread(image_path).shape[:2]
ET.SubElement(size, "width").text = str(img_width)
ET.SubElement(size, "height").text = str(img_height)
# 添加物体信息
obj = ET.SubElement(annotation, "object")
ET.SubElement(obj, "name").text = calss_name
ET.SubElement(obj, "pose").text = "Unspecified"
ET.SubElement(obj, "truncated").text = "0"
ET.SubElement(obj, "difficult").text = "0"
bndbox = ET.SubElement(obj, "bndbox")
ET.SubElement(bndbox, "xmin").text = str(xmin)
ET.SubElement(bndbox, "ymin").text = str(ymin)
ET.SubElement(bndbox, "xmax").text = str(xmax)
ET.SubElement(bndbox, "ymax").text = str(ymax)
# 保存 XML 文件
tree = ET.ElementTree(annotation)
tree.write(xml_path, encoding='utf-8', xml_declaration=True)
else:
# 如果该帧的 XML 文件已存在,则只需添加新的物体信息
tree = ET.parse(xml_path)
root = tree.getroot()
obj = ET.SubElement(root, "object")
ET.SubElement(obj, "name").text = calss_name
ET.SubElement(obj, "pose").text = "Unspecified"
ET.SubElement(obj, "truncated").text = "0"
ET.SubElement(obj, "difficult").text = "0"
bndbox = ET.SubElement(obj, "bndbox")
ET.SubElement(bndbox, "xmin").text = str(xmin)
ET.SubElement(bndbox, "ymin").text = str(ymin)
ET.SubElement(bndbox, "xmax").text = str(xmax)
ET.SubElement(bndbox, "ymax").text = str(ymax)
# 保存更新后的 XML 文件
tree.write(xml_path, encoding='utf-8', xml_declaration=True)
if __name__ == "__main__":
uavdtmot_gt_path = 'your/path/to/UAVDT-MOT/labels/GT'
images_path = 'your/path/to/UAVDT-MOT/images'
voc_path = 'your/path/to/UAVDT-MOT/Annotations'
class_names = ['car', 'truck', 'bus']
if not os.path.exists(uavdtmot_voc_path):
os.makedirs(uavdtmot_voc_path)
process_files(uavdtmot_gt_path, images_path, voc_path , class_names)
UAVDT数据集的主要任务是多目标跟踪,也可以兼顾目标检测任务从官网上下载图片和标注文件UAVDT。也可以直接点这个下载链接UAVDT下载链接。
下载完后会得到两个文件
UAV-benchmark-M.zip就是图片文件夹,UAV-benchmark-MOTD_v1.0.zip下边有一个GT文件夹,这个文件夹就是标注文件夹。UAVDT提供三种标注形式,需要转成目标检测数据集则只用关心_gt_whole.txt文件。
根据READEME文件显示_gt_whole.txt文件的格式为
这些已经包含了目标检测数据集所需要的全部信息了,只要做好行和图片文件的对应,转成VOC格式就手拿把掐了。