DOTA转YOLO

最新推荐文章于 2024-04-12 10:52:27 发布

Oct _

最新推荐文章于 2024-04-12 10:52:27 发布

阅读量1.5k

点赞数 2

文章标签： python 深度学习开发语言

本文链接：https://blog.csdn.net/weixin_49044667/article/details/127053716

版权

#DOTA_to_YOLO.py:

import data_util as util
import os
import numpy as np
from PIL import Image

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
Image.MAX_IMAGE_PIXELS = None


## trans dota format to format YOLO(darknet) required
def dota2darknet(imgpath, txtpath, dstpath, extractclassname):
    """
    :param imgpath: the path of images
    :param txtpath: the path of txt in dota format
    :param dstpath: the path of txt in YOLO format
    :param extractclassname: the category you selected
    :return:
    """
    filelist = util.GetFileFromThisRootDir(txtpath)
    for fullname in filelist:
        objects = util.parse_dota_poly(fullname)
        name = os.path.splitext(os.path.basename(fullname))[0]
        img_fullname = os.path.join(imgpath, name + '.png')
        img = Image.open(img_fullname)
        img_w, img_h = img.size
        # print img_w,img_h
        with open(os.path.join(dstpath, name + '.txt'), 'w') as f_out:
            for obj in objects:
                poly = obj['poly']
                bbox = np.array(util.dots4ToRecC(poly, img_w, img_h))
                if (sum(bbox <= 0) + sum(bbox >= 1)) >= 1:
                    continue
                if (obj['name'] in extractclassname):
                    id = extractclassname.index(obj['name'])
                else:
                    continue
                outline = str(id) + ' ' + ' '.join(list(map(str, bbox)))
                f_out.write(outline + '\n')


if __name__ == '__main__':
    dota2darknet('E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\train1.5\images',
                 'E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\train1.5\labels',
                 'E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\train1.5\labels_yolo',
                 util.wordname_18)

    dota2darknet('E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\val1.5\images',
                 'E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\val1.5\labels',
                 'E:\BaiduNetdiskDownload\DOTA1.5\DOTA\\val1.5\labels_yolo',
                 util.wordname_18)

data_util.py

import sys
import codecs
import numpy as np
import shapely.geometry as shgeo
import os
import re
import math
"""
    some basic functions which are useful for process DOT