将实例分割数据集转为目标检测数据集

版权声明：转载请注明出处，谢谢合作 https://blog.csdn.net/qq_34199326/article/details/83819140

大家都知道的，想要让深度学习算法发挥良好的性能，数据集的好坏至关重要。如果做目标检测，有时候苦于想做的项目没有合适的数据集，而分割数据集有自己项目课题需要的，或者就是想把高质量的分割数据集（如Cityscapes）拿来用于目标检测任务，就需要将深度学习实例分割数据集改造为目标检测数据集。

那么如何将像素级的多边形标注的分割数据标注转为目标检测的bbox标注呢？其实很简单，由于分割的每个目标的多边形标注都是标注多个像素[x,y]从而标注出多边形形状，只需要找出这个目标的多边形像素标注的Ｘmin,Ymin,Xmax,Ymax即可得到包围框。下面以CityScapes数据集转为目标检测数据集（这里转为适合YOLO系列需要的输入格式.txt）为例给出代码（实际上其他分割数据集均适用）。YOLO所需的格式如下图：

<object-class> <x> <y> <width> <height>

 代码如下：（默认一张图片对应一个.json标签）

import json
import os 
from os import listdir, getcwd
from os.path import join
import os.path
rootdir='/home/wang/下载/数据集/cityscapes/Images/train'　#写自己存放图片的数据地址
def position(pos):　#该函数用来找出xmin,ymin,xmax,ymax即bbox包围框
    x=[]
    y=[]
    nums=len(pos)
    for i in range(nums):
        x.append(pos[i][0])
        y.append(pos[i][1])
    x_max=max(x)
    x_min=min(x)
    y_max=max(y)
    y_min=min(y)
    b=(float(x_min),float(x_max),float(y_min),float(y_max))
    return b

def convert(size, box):　#该函数将xmin,ymin,xmax,ymax转为x,y,w,h中心点坐标和宽高
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(image_id):   
    load_f=open("./train/%s.json"%(image_id),'r')#导入json标签的地址
    load_dict = json.load(load_f)
    out_file = open('./voc_type/train/%s.txt'%(image_id), 'w') #输出标签的地址
    #keys=tuple(load_dict.keys()) 
    w=load_dict['imgWidth']  #原图的宽，用于归一化
    h=load_dict['imgHeight']
    #print(h)
    objects=load_dict['objects']
    nums=len(objects)
    #print(nums)
    #object_key=tuple(objects.keys()
   
    
    
    for i in range(0,nums):
        labels=objects[i]['label']
        #print(i)
        if (labels in ['person','rider']):
            #print(labels)
            pos=objects[i]['polygon']   
            b=position(pos)
            bb = convert((w,h), b)
            cls_id=2  #我这里把行人和骑自行车的人都设为类别２
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
            #print(type(pos))
        elif (labels in ['car','truck','bus','caravan','trailer']):
            #print(labels)
            pos=objects[i]['polygon']
            b=position(pos)
            bb = convert((w,h), b)
            cls_id=1 #我这里把各种类型的车都设为类别１
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

def image_id(rootdir):
    a=[]
    for parent,dirnames,filenames in os.walk(rootdir):
        for filename in filenames:
            filename=filename.strip('.png')
            #print(filename)
            a.append(filename)
    return a
names=image_id(rootdir)
for image_id in names:
    convert_annotation(image_id)

代码只需要简单修改就可以用于各种任务中去。

当然，有时候图片的数目与标签数目不一致该咋办？（例如有的图片没标注信息或者没有.json文件）

此时只需要将图片文件夹与标签文件夹进行集合的求补运算，然后删去这些图片即可，代码如下：

#encoding:utf-8
import os
import os.path
from os import listdir, getcwd
from os.path import join
txt_dir='./label/train'   #你的图片文件夹路径
pic_dir='./Image/train'　#你的标签文件夹路径


def txt(rootdir):
    a=[]
    for parent,dirnames,filenames in os.walk(rootdir):
        for filenames in filenames:
            filenames=filenames.strip('.txt') #这里看你的标注文件后缀是什么就改为什么
            a.append(filenames)
    return a
def pic(rootdir):
    b=[]
    for parent,dirnames,filenames in os.walk(rootdir):
        for filenames in filenames:
            filenames=filenames.strip('.jpg')　#图片的后缀是什么就改为什么
            b.append(filenames)
    return b

txt_set=txt(txt_dir)
txt_set=set(txt_set)
pic_set=pic(pic_dir)
pic_set=set(pic_set)
#comp=txt_set-pic_set
comp=pic_set-txt_set　#图片比标注多时，进行做差
print("ok")
print(len(comp))

for item in comp:　#删去这些无标注的图片
    file=pic_dir+'/'+item+'.jpg'
    if os.path.exists(file):
        os.remove(file)
        print(file)
        
#for item in comp:
 #   file=txt_dir+'/'+item+'.json'
  #  if os.path.exists(file):
   #     os.remove(file)
    #    print(file)

接下来，还有个问题，如果标注信息是所有的图片的标注信息都写在一个.json文件中咋办（类似于BDDV数据集和coco数据集那样）。其实也简单，代码如下（BDDV数据集有bbox信息，就直接提取了，如果也是分割数据，那么就如之前程序那样从多边形中提取bbox包围框即可）

import json
import os

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

load_f=open('bdd100k_labels_images_train.json','r')
data=json.load(load_f)
#print(data[0].keys())

w=1280
h=720

if not os.path.exists('./labels_yolo'):
    os.makedirs('./labels_yolo')


num=len(data)
for i in range(num):
    pic_name=data[i]['name']
    name=pic_name.split('.')[0]
    labels=data[i]['labels']
    obj_num=len(labels)
    for j in range(obj_num):
        category=labels[j]['category']
        if (category in ['bus','truck','car']):
            out_file=open('./labels_yolo/%s.txt'%(name),'a')
            cls_id=1
            pos=labels[j]['box2d']
            x_min=pos['x1']
            x_max=pos['x2']
            y_min=pos['y1']
            y_max=pos['y2']
            b = (float(x_min), float(x_max), float(y_min), float(y_max))
            bb = convert((w,h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
        if (category in ['person','rider']):
            out_file=open('./labels_yolo/%s.txt'%(name),'a')
            cls_id=2
            pos=labels[j]['box2d']
            x_min=pos['x1']
            x_max=pos['x2']
            y_min=pos['y1']
            y_max=pos['y2']
            b = (float(x_min), float(x_max), float(y_min), float(y_max))
            bb = convert((w,h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')