目标检测对输入数据的格式要求不尽相同,本文提供的是yolo转VOC格式的代码。
有时候yolo的标注文件比较乱,特别是图片和标注文件并非完全对应,这时需要先做一个处理。我们利用交集的方法,来保证图片和标签的一一对应。
代码如下:
imgs=sorted(os.listdir('/yourpath/images'))
lbs=sorted(os.listdir('/yourpath/labels'))
l = [os.path.splitext(i)[0] for i in lbs] #去掉图片的后缀
b = [os.path.splitext(i)[0] for i in imgs] # 去掉标注文件的后缀
a=list(set(b) & set(l)) # 这就得到图片和标注文件的一一对应
接着,把一一对应的图片和标签文件取出来,从而得到一个标准的正确的yolo目标检测数据集,代码如下:
for i in tqdm(a):
img = os.path.join('/yourpath/images/',i+'.jpg')
lb = os.path.join('/yourpath/labels/',i+'.txt')
shutil.copy(img,'/yournewpath/images')
shutil.copy(lb,'/yournewpath/labels')
直接上代码,相关说明看代码的注释
import os, cv2, shutil
from lxml import etree, objectify
from tqdm import tqdm
from PIL import Image
import numpy as np
import time
import os.path as osp
import json
def cover_copy(src,dst):
'''
src和dst都必须是文件,该函数是执行覆盖操作
'''
if os.path.exists(dst):
os.remove(dst)
shutil.copy(src,dst)
else:
shutil.copy(src,dst)
def yolo2voc(sourcedir,savedir,class_names):
"""_summary_
Args:
sourcedir (_type_): /yournewpath 写到images 和labels上一级
savedir (_type_): 写到JPEGImages和Annotations上一级
class_names(list): yolo数据中所有的类别,顺序要与索引对应上
"""
img_savepath= osp.join(savedir,'JPEGImages')
ann_savepath=osp.join(savedir,'Annotations')
for p in [img_savepath,ann_savepath]:
if osp.exists(p):
shutil.rmtree(p)
os.makedirs(p)
else:
os.makedirs(p)
filenames = os.listdir(osp.join(sourcedir,'images'))
for filename in tqdm(filenames):
filepath = osp.join(sourcedir,'images',filename)
annfilepath=osp.join(sourcedir,'labels',osp.splitext(filename)[0]+'.txt')
annopath = osp.join(ann_savepath,osp.splitext(filename)[0] + ".xml") #生成的xml文件保存路径
dst_path = osp.join(img_savepath,filename)
im = Image.open(filepath)
image = np.array(im).astype(np.uint8)
w=image.shape[1]
h=image.shape[0]
cover_copy(filepath, dst_path)#把原始图像复制到目标文件夹
anns=[i.strip().split() for i in open(annfilepath).readlines()]
objs = []
for ann in anns:
name = class_names[int(ann[0])]
xcenter = float(ann[1])*w
ycenter = float(ann[2])*h
bw=float(ann[3])*w
bh=float(ann[4])*h
xmin = (int)(xcenter-bw/2)
ymin = (int)(ycenter-bh/2)
xmax = (int)(xcenter+bw/2)
ymax = (int)(ycenter+bh/2)
obj = [name, 1.0, xmin, ymin, xmax, ymax]
#标错框在这里
if not(xmin-xmax==0 or ymin-ymax==0):
objs.append(obj)
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder('VOC'),
E.filename(filename),
E.source(
E.database('YOLO'),
E.annotation('VOC'),
E.image('YOLO')
),
E.size(
E.width(image.shape[1]),
E.height(image.shape[0]),
E.depth(image.shape[2])
),
E.segmented(0)
)
for obj in objs:
E2 = objectify.ElementMaker(annotate=False)
anno_tree2 = E2.object(
E.name(obj[0]),
E.pose(),
E.truncated("0"),
E.difficult(0),
E.bndbox(
E.xmin(obj[2]),
E.ymin(obj[3]),
E.xmax(obj[4]),
E.ymax(obj[5])
)
)
anno_tree.append(anno_tree2)
etree.ElementTree(anno_tree).write(annopath, pretty_print=True)