DOTA数据集的处理(图片切割,转换为YOLO格式全流程)

本人最近下载了DOTA1.5版本的数据集,在此将处理的流程在此与大家做一个分享:

目的:将DOTA数据集处理为YOLO模型可以进行训练的形式

在拿到手的DOTA数据集标注样式如下:

7c72db1683dc4644b6cff2b03ee5696d.png

其各个部分的含义为:

imagesource: GoogleEarth:这表示图像的来源是Google Earth

gsd: 0.146343590398:GSD(Ground Sample Distance)表示地面采样距离,即图像中每个像素所代表的实际地面面积的大小。这里的0.146343590398米/像素意味着图像中的一个像素对应地面上大约0.146米的距离。

937 913 921 912 923 874 940 875 small-vehicle 0:前八个数代表从目标框左上角开始的(x,y)坐标,沿顺时针方向的四点坐标,small-vehicle为类别,后面紧跟的一个数字代表检测的难易程度。

我们首先去除每个标签文件用不上的前两行,代码如下:

import os

def remove_first_two_lines(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
    with open(file_path, 'w', encoding='utf-8') as file:
        file.writelines(lines[2:])

directory = 'E:/UserTable/DOTA1.5/val/labelTxt-v1.5'  # 指定要处理的文件夹路径,这里使用当前文件夹
for filename in os.listdir(directory):
    if filename.endswith('.txt'):
        file_path = os.path.join(directory, filename)
        remove_first_two_lines(file_path)

去除了前两行代码,我们可以使用一下代码对图片和对应标签进行分割:

import cv2
import os

#  图像宽不足裁剪宽度,填充至裁剪宽度
def fill_right(img, size_w):
    size = img.shape
    #  填充值为数据集均值
    img_fill_right = cv2.copyMakeBorder(img, 0, 0, 0, size_w - size[1],
                                        cv2.BORDER_CONSTANT, value = (107, 113, 115))
    return img_fill_right

#  图像高不足裁剪高度,填充至裁剪高度
def fill_bottom(img, size_h):
    size = img.shape
    img_fill_bottom = cv2.copyMakeBorder(img, 0, size_h - size[0], 0, 0,
                                         cv2.BORDER_CONSTANT, value = (107, 113, 115))
    return img_fill_bottom

#  图像宽高不足裁剪宽高度,填充至裁剪宽高度
def fill_right_bottom(img, size_w, size_h):
    size = img.shape
    img_fill_right_bottom = cv2.copyMakeBorder(img, 0, size_h - size[0], 0, size_w - size[1],
                                               cv2.BORDER_CONSTANT, value = (107, 113, 115))
    return img_fill_right_bottom

#  图像切割
#  img_floder 图像文件夹
#  out_img_floder 图像切割输出文件夹
#  size_w 切割图像宽
#  size_h 切割图像高
#  step 切割步长
def image_split(img_floder, out_img_floder, size_w = 1000, size_h = 1000, step = 800):
    img_list = os.listdir(img_floder)
    count = 0
    for img_name in img_list:
        number = 0
        #  去除.png后缀
        name = img_name[:-4]
        img = cv2.imread(img_floder + "\\" + img_name)
        size = img.shape
        #  若图像宽高大于切割宽高
        if size[0] >= size_h and size[1] >= size_w:
           count = count + 1
           for h in range(0, size[0] - 1, step):
               start_h = h
               for w in range(0, size[1] - 1, step):
                   start_w = w
                   end_h = start_h + size_h
                   if end_h > size[0]:
                      start_h = size[0] - size_h
                      end_h = start_h + size_h
                   end_w = start_w + size_w
                   if end_w > size[1]:
                      start_w = size[1] - size_w
                   end_w = start_w + size_w
                   cropped = img[start_h : end_h, start_w : end_w]
                   #  用起始坐标来命名切割得到的图像,为的是方便后续标签数据抓取
                   name_img = name + '_'+ str(start_h) +'_' + str(start_w)
                   cv2.imwrite('{}/{}.png'.format(out_img_floder, name_img), cropped)
                   number = number + 1
        #  若图像高大于切割高,但宽小于切割宽
        elif size[0] >= size_h and size[1] < size_w:
            print('图片{}需要在右面补齐'.format(name))
            count = count + 1
            img0 = fill_right(img, size_w)
            for h in range(0, size[0] - 1, step):
               start_h = h
               start_w = 0
               end_h = start_h + size_h
               if end_h > size[0]:
                  start_h = size[0] - size_h
                  end_h = start_h + size_h
               end_w = start_w + size_w
               cropped = img0[start_h : end_h, start_w : end_w]
               name_img = name + '_' + str(start_h) + '_' + str(start_w)
               cv2.imwrite('{}/{}.png'.format(out_img_floder, name_img), cropped)
               number = number + 1
        #  若图像宽大于切割宽,但高小于切割高
        elif size[0] < size_h and size[1] >= size_w:
            count = count + 1
            print('图片{}需要在下面补齐'.format(name))
            img0 = fill_bottom(img, size_h)
            for w in range(0, size[1] - 1, step):
               start_h = 0
               start_w = w
               end_w = start_w + size_w
               if end_w > size[1]:
                  start_w = size[1] - size_w
                  end_w = start_w + size_w
               end_h = start_h + size_h
               cropped = img0[start_h : end_h, start_w : end_w]
               name_img = name + '_'+ str(start_h) +'_' + str(start_w)
               cv2.imwrite('{}/{}.png'.format(out_img_floder, name_img), cropped)
               number = number + 1
        #  若图像宽高小于切割宽高
        elif size[0] < size_h and size[1] < size_w:
            count = count + 1
            print('图片{}需要在下面和右面补齐'.format(name))
            img0 = fill_right_bottom(img,  size_w, size_h)
            cropped = img0[0 : size_h, 0 : size_w]
            name_img = name + '_'+ '0' +'_' + '0'
            cv2.imwrite('{}/{}.png'.format(out_img_floder, name_img), cropped)
            number = number + 1
        print('{}.png切割成{}张.'.format(name,number))
    print('共完成{}张图片'.format(count))

#  txt切割
#  out_img_floder 图像切割输出文件夹
#  txt_floder txt文件夹
#  out_txt_floder txt切割输出文件夹
#  size_w 切割图像宽
#  size_h 切割图像高
def txt_split(out_img_floder, txt_floder, out_txt_floder, size_h = 1000, size_w = 1000):
    img_list = os.listdir(out_img_floder)
    for img_name in img_list:
        #  去除.png后缀
        name = img_name[:-4]
        #  得到原图像(也即txt)索引 + 切割高 + 切割宽
        name_list = name.split('_')
        txt_name = name_list[0]
        h = int(name_list[1])
        w = int(name_list[2])
        txtpath = txt_floder + "\\" + txt_name + '.txt'
        out_txt_path = out_txt_floder + "\\" + name + '.txt'
        f = open(out_txt_path, 'a')
        #  打开txt文件
        i=0
        with open(txtpath, 'r') as f_in:
             lines = f_in.readlines()
             #  逐行读取
             for line  in lines:
                if i<2:
                    i = i+1
                else:
                     splitline = line.split(' ')
                     label = splitline[8]
                     difficult = splitline[9]
                     x1 = int(float(splitline[0]))
                     y1 = int(float(splitline[1]))
                     x2 = int(float(splitline[2]))
                     y2 = int(float(splitline[3]))
                     x3 = int(float(splitline[4]))
                     y3 = int(float(splitline[5]))
                     x4 = int(float(splitline[6]))
                     y4 = int(float(splitline[7]))
                     if w <= x1 <= w + size_w and w <= x2 <= w + size_w and \
                     w <= x3 <= w + size_w and w <= x4 <= w + size_w and \
                     h <= y1 <= h + size_h and h <= y2 <= h + size_h and \
                     h <= y3 <= h + size_h and h <= y4 <= h + size_h:
                         f.write('{} {} {} {} {} {} {} {} {} {}'.format(int(x1 - w),
                                 int(y1 - h), int(x2 - w), int(y2 - h), int(x3 - w),
                                 int(y3 - h), int(x4 - w), int(y4 - h),
                                 label, difficult))
        f.close()
        print('{}.txt切割完成.'.format(name))

#  图像数据集文件夹
img_floder = r'E:/UserTable/DOTA1.5/val/images/images'
#  切割得到的图像数据集存放文件夹
out_img_floder = r'E:/UserTable/DOTA1.5/val/images/images_qiege'
#  txt数据集文件夹
txt_floder = r'E:/UserTable/DOTA1.5/val/labelTxt-v1.5'
#  切割后数据集的标签文件存放文件夹
out_txt_floder = r'E:/UserTable/DOTA1.5/val/labels_qiege'
#  切割图像宽
size_w = 1024
#  切割图像高
size_h = 1024
#  切割步长,重叠度为size_w - step
step = 824

image_split(img_floder, out_img_floder, size_w, size_h, step)
txt_split(out_img_floder, txt_floder, out_txt_floder, size_h, size_w)

该部分代码引自:DOTA数据集切割,图像大小1024x1024,切割步长824,重叠度200

最后,使用以下代码完成YOLO标注形式的数据集标签转换:

import os
from PIL import Image

def convert_dota_to_yolo(dota_folder):
    yolo_folder = "E:/UserTable/DOTA1.5/val/labels_yolo"  # 存储转换后的yolo格式文件夹
    made = "E:/UserTable/DOTA1.5/val/images/images_qiege"
    os.makedirs(yolo_folder, exist_ok=True)
    label_map = {'airport':0,
  'small-vehicle':1,
  'large-vehicle':2,
  'plane':3,
  'storage-tank':4,
  'ship':5,
  'harbor':6,
  'ground-track-field':7,
  'soccer-ball-field':8,
  'tennis-court':9,
  'swimming-pool':10,
  'baseball-diamond':11,
  'roundabout':12,
  'basketball-court':13,
  'bridge':14,
  'helicopter':15,
  'container-crane':16,
  'helipad':17}
    label_index = 1
    
    for filename in os.listdir(dota_folder):
        if filename.endswith(".txt"):
            dota_file = os.path.join(dota_folder, filename)
            yolo_file = os.path.join(yolo_folder, filename.replace(".txt", ".txt"))
            caonimade = os.path.join(made, filename.replace(".txt", ".png"))
            
            with open(dota_file, "r") as dota_f:
                with open(yolo_file, "w") as yolo_f:
                    for line in dota_f:
                        line = line.strip().split()
                        class_name = line[8]
                        x1, y1, x2, y2, x3, y3, x4, y4 = map(float, line[:8])
                        
                        # 如果类别名称不在字典中,则添加新的索引
                        if class_name not in label_map:
                            raise ValueError("有未命名的东西,giao!!!")
                        
                        width, height = get_image_size(caonimade)

                        # 计算yolo格式的中心点坐标和宽高
                        x_center = (x1 + x3) / 2 / width
                        y_center = (y1 + y3) / 2 / height
                        width = (x3 - x1) / width
                        height = (y3 - y1) / height

                        # 将转换后的结果写入yolo格式文件
                        yolo_f.write(f"{label_map[class_name]} {x_center} {y_center} {width} {height}\n")
    
    print("转换完成!")

def get_image_size(image_path):
    img = Image.open(image_path)
    width, height = img.size
    return width, height

# 调用函数进行转换
convert_dota_to_yolo("E:/UserTable/DOTA1.5/val/labels_qiege")

 就此完成处理的全部流程!

完整代码同样打包在了:https://download.csdn.net/download/wbk1669177085/89888723,可以免费下载

评论 14
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值