dota数据集标注改coco标注

#这是为了把dota类型的标注转为coco的正框标注
#原因是要做一个车辆检测，coco格式标注，但是dota更适合用来训练
#下面是代码

#一些模块为了后面的可视化而增加
import os.path as osp
import os
from PIL import Image
import json
from matplotlib import pyplot as plt

#根据信息生成模板，此处要求bbox必须是xyxy标注
def coco_annotations(bbox,cid,bbox_id,img_id,iscrowd):
    x1,y1,x2,y2=bbox
    return {'segmentation':[[x1,y1,x2,y1,x2,y2,x1,y2]],
           'bbox':[x1,y1,x2-x1+1,y2-y1+1],
           'category_id':cid,
           'area':(y2-y1+1)*(x2-x1+1),
           'iscrowd':iscrowd,
           'image_id':img_id,
           'id':bbox_id}
def coco_images(file_name,height,width,img_id):
    return {'file_name':file_name,
           'height':height,
           'width':width,
           'id':img_id}

#从dota数据集中的txt格式标注中获取信息，函数中需要用到变量cls_name2id，categories
def deal_with_txt(label_path, img_id, anno_id):
    annos = []
    
    with open(label_path,'r') as gt:
        gt_lines=gt.readlines()
    for i in gt_lines:
        iscrowd=int(i[-2])
        i=i.split(' ')
        cls_name=i[-2]
        cid=cls_name2id[cls_name]
        
        cood=i[:8]
        cood=tuple(map(float,cood))
        x1=min(cood[0],cood[2],cood[4],cood[6])
        x2=max(cood[0],cood[2],cood[4],cood[6])
        y1=min(cood[1],cood[3],cood[5],cood[7])
        y2=max(cood[1],cood[3],cood[5],cood[7])
        b=(x1,y1,x2,y2)
        #下面使用coco_annotations函数把上面的读取数据输入
        anno=coco_annotations(b,cid,anno_id,img_id,iscrowd)
        annos.append(anno)
        anno_id+=1
    return annos,anno_id


def generate_coco_fmt(data_root,anno_root,categories,img_root):
    '''
    data_root是数据集路径
    下面的路径是相对data_root的路径
    anno_root是数据集的标注文件的路径
    img_root是数据集的图片路径
    
    categories需要事先得到
    '''
    img_id, anno_id = 0, 0
    all_annos, all_images = [], []
    
    for anno_txt in os.listdir(osp.join(data_root,anno_root)):
        file_name=anno_txt.replace('txt','png')
        label_path=osp.join(data_root,anno_root,anno_txt)
        img_path=osp.join(data_root,img_root,file_name)
        if osp.exists(img_path):
            annos,anno_id=deal_with_txt(label_path,img_id,anno_id)
            all_annos.extend(annos)
            
            img=Image.open(img_path)
            all_images.append(coco_images(osp.join(img_root,file_name),img.height,img.width,img_id))
            img_id+=1
    return{
        'images':all_images,
        "annotations": all_annos,
        "categories": categories,
        "type": "instance"
    }

'''
	这一部分的目的是为了获取必须提前获取的categories和cls_name2id
'''
classes=('plane', 'baseball-diamond', 'bridge', 'ground-track-field',
               'small-vehicle', 'large-vehicle', 'ship', 'tennis-court',
               'basketball-court', 'storage-tank', 'soccer-ball-field',
               'roundabout', 'harbor', 'swimming-pool', 'helicopter')
cls_name2id={}
categories=[]
for i in range(1,16):
    cls_name2id.update({classes[i-1]:i})
    dict_cate={'id':i,'name':classes[i-1],'supercategory':classes[i-1]}
    categories.append(dict_cate)

#路径自己改一下记得
data_root='../data/DOTA1-split-1024'
anno_root='trainval1024/annfiles'
img_root='trainval1024/images'
dota_coco_fmt=generate_coco_fmt(data_root,anno_root,categories,img_root)
json.dump(dota_coco_fmt,open(osp.join(data_root,'train.json'),'a'),ensure_ascii=False)

效果图：
原始的txt的包含信息：
[‘597.0 429.0 614.0 432.0 608.0 473.0 591.0 470.0 small-vehicle 1\n’,
‘409.0 647.0 410.0 667.0 343.0 677.0 342.0 658.0 small-vehicle 0\n’,
‘434.0 959.0 505.0 955.0 507.0 978.0 436.0 980.0 large-vehicle 0\n’,
‘697.0 757.0 718.0 754.0 743.0 881.0 721.0 886.0 large-vehicle 0\n’,
‘556.0 790.0 580.0 786.0 603.0 913.0 578.0 918.0 large-vehicle 0\n’,
‘521.0 772.0 544.0 767.0 549.0 837.0 528.0 840.0 small-vehicle 0\n’,
‘456.0 790.0 479.0 789.0 505.0 939.0 483.0 940.0 large-vehicle 0\n’,
‘480.0 781.0 503.0 773.0 539.0 929.0 509.0 941.0 large-vehicle 0\n’,
‘607.0 729.0 633.0 723.0 680.0 898.0 647.0 909.0 large-vehicle 0\n’,
‘390.0 360.0 412.0 364.0 409.0 449.0 382.0 446.0 large-vehicle 0\n’,
‘751.0 21.0 747.0 43.0 584.0 0.0 669.0 0.0 large-vehicle 0\n’,]
以及可视化后的图片：

当然，由于coco的格式是所有标注统一在一起的json文件，这里没法展示：
于是展示一下，正框可视化后的结果。
如果想知道可视化代码可以去我的历史文章里面看。
如果自己写可视化的话，使用opencv比Image模块要方便很多，Image模块需要numpy之间转来转去，opencv有很多内置函数，记不得了上网搜索很方便，而且opencv还支持标注框附带文字。