最近在搞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改成自己的路径就行。