Kaggle-Airbus-Ships-Detection数据集转换Dota格式

最近在搞OBB目标检测，发现旋转框数据集很是稀有，找半天发现kaggle上的这个数据集不错，想要使用mmrotate框架检测的话就要dota数据集的格式，于是简单实现了一下将kaggle airbus的数据集标注格式(csv)转换成dota格式(txt)

kaggle数据集网站：https://www.kaggle.com/competitions/airbus-ship-detection

dota数据集网站：DOTA

import pandas as pd
import numpy as np
from tqdm import tqdm
from rasterio import features
from shapely.geometry import MultiPoint


# function to convert RLE mask into 2d pixel array
def encode_mask(mask, shape=(768,768)):
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    s = mask.split()
    for i in range(len(s)//2):
        start = int(s[2*i]) - 1
        length = int(s[2*i+1])
        img[start:start+length] = 1
    return np.transpose(img.reshape(shape))


def main():
    RLE_PATH = '../datasets/train_ship_segmentations_v2.csv'
    FILE_PATH = '../datasets/txt/'
    # read the RLE annotation in a Pandas DataFrame
    rle_df = pd.read_csv(RLE_PATH)
    rle_df.head()
    processed_image = set()
    for row in tqdm(rle_df.itertuples(), total=len(rle_df)):
        if pd.isna(row.EncodedPixels):
            continue
        del_part = '.jpg'
        FILE_NAME = f'{FILE_PATH}{row.ImageId}'
        FILE_NAME = FILE_NAME.replace(del_part, '')
        if row.ImageId not in processed_image:
            with open(f'{FILE_NAME}.txt', 'w') as file:
                file.write('imagesource:Kaggle\n')
                file.write('gsd:NaN\n')
                processed_image.add(row.ImageId)
        if row.ImageId in processed_image:
            # convert RLE encoded pixels into binary mask
            mask = encode_mask(str(row.EncodedPixels), shape=(768, 768))
            # vectorize mask into GeoJSON
            value = 0.0
            for polygon, value in list(features.shapes(mask)):
                if value == 1.0:
                    break
            if value != 1.0:
                print('Error while vectorizing mask')

            # get oriented bounding box around shape
            coords = polygon['coordinates'][0]
            obbox = MultiPoint(coords).minimum_rotated_rectangle

            # get external coordinates of oriented rectangle
            # compute length, width and angle
            p1, p2, p3, p4, p5 = list((obbox.exterior.coords))
            with open(f'{FILE_NAME}.txt', 'a') as file:
                file.write(f'{int(p1[0])} {int(p1[1])} {int(p2[0])} {int(p2[1])} {int(p3[0])} {int(p3[1])} {int(p4[0])} {int(p4[1])} ship 0\n')


if __name__ == '__main__':
    main()

将代码中的RLE_PATH和FILE_PATH改成自己的路径就行。