1.下载coco数据集
(1)通过官网下载。数据集很大,总是会中断,还要重头再来。。
http://cocodataset.org/#overview
(2)YOLO官网给出的方法(脚本是在darknet-master文件里的~)
https://pjreddie.com/darknet/yolo/
cp scripts/get_coco_dataset.sh data
cd data
bash get_coco_dataset.sh
之后就下载了,不怕中断,可以断点继续。下载的是train2014(14G)和val2014(2.5G)。
2.提取coco中部分类别并转换成voc格式
先上代码
from pycocotools.coco import COCO
import skimage.io as io
import matplotlib.pyplot as plt
import pylab, os, cv2, shutil
from lxml import etree, objectify
from tqdm import tqdm
import random
from PIL import Image
pylab.rcParams['figure.figsize'] = (8.0, 10.0)
dataDir = '/xxx/bs_darknet-master/data/coco/images' #train2014文件夹所在路径(train2014存放的是coco所有的图)
CK5cats = ['tie','person']#想要提取的类别名称
CKdir = "/xxx/bs_darknet-master/data/coco"
CKimg_dir = CKdir + "/" + "images/train"#提取出的类别图片存储路径
CKanno_dir = CKdir + "/" + "images/xml"#对应生成的xml存储路径
def mkr(dir):
if not os.path.exists(dir):
os.makedirs(dir)
def showimg(coco, dataType, img, CK5Ids):
global dataDir
I = io.imread('%s/%s/%s' % (dataDir, dataType, img['file_name']))
plt.imshow(I)
plt.axis('off')
annIds = coco.getAnnIds(imgIds=img['id'], catIds=CK5Ids, iscrowd=None)
anns = coco.loadAnns(annIds)
coco.showAnns(anns)
plt.show()
def save_annotations(dataType, filename, objs):
annopath = CKanno_dir + "/" + filename[:-3] + "xml"
img_path = dataDir + "/" + dataType + "/" + filename
dst_path = CKimg_dir + "/" + filename
img = cv2.imread(img_path)
im = Image.open(img_path)
if im.mode != "RGB":
print(filename + " not a RGB image")
im.close()
return
im.close()
shutil.copy(img_path, dst_path)
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder('1'),
E.filename(filename),
E.source(
E.database('CKdemo'),
E.annotation('VOC'),
E.image('CK')
),
E.size(
E.width(img.shape[1]),
E.height(img.shape[0]),
E.depth(img.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)
def showbycv(coco, dataType, img, classes, CK5Ids):
global dataDir
filename = img['file_name']
filepath = '%s/%s/%s' % (dataDir, dataType, filename)
I = cv2.imread(filepath)
annIds = coco.getAnnIds(imgIds=img['id'], catIds=CK5Ids, iscrowd=None)
anns = coco.loadAnns(annIds)
objs = []
for ann in anns:
name = classes[ann['category_id']]
if name in CK5cats:
if 'bbox' in ann:
bbox = ann['bbox']
xmin = (int)(bbox[0])
ymin = (int)(bbox[1])
xmax = (int)(bbox[2] + bbox[0])
ymax = (int)(bbox[3] + bbox[1])
obj = [name, 1.0, xmin, ymin, xmax, ymax]
objs.append(obj)
cv2.rectangle(I, (xmin, ymin), (xmax, ymax), (255, 0, 0))
cv2.putText(I, name, (xmin, ymin), 3, 1, (0, 0, 255))
save_annotations(dataType, filename, objs)
cv2.imshow("img", I)
cv2.waitKey(1)
def catid2name(coco):
classes = dict()
for cat in coco.dataset['categories']:
classes[cat['id']] = cat['name']
# print(str(cat['id'])+":"+cat['name'])
return classes
def get_CK5():
mkr(CKimg_dir)
mkr(CKanno_dir)
dataTypes = ['train2014']
for dataType in dataTypes:
annFile = '/xxx/bs_darknet-master/data/coco/annotations/instances_train2014.json'.format(dataDir, dataType) #这里改成自己的instances_train2014.json路径
coco = COCO(annFile)
CK5Ids = coco.getCatIds(catNms=CK5cats)
classes = catid2name(coco)
for srccat in CK5cats:
print(dataType + ":" + srccat)
catIds = coco.getCatIds(catNms=[srccat])
imgIds = coco.getImgIds(catIds=catIds)
# imgIds=imgIds[0:100]
for imgId in tqdm(imgIds):
img = coco.loadImgs(imgId)[0]
showbycv(coco, dataType, img, classes, CK5Ids)
# showimg(coco,dataType,img,CK5Ids)
if __name__ == "__main__":
get_CK5()
将脚本放入和train2014文件夹同目录下,运行即可。会生成两个文件夹,分别存放提取的图片和对应xml标签。
参考资料
https://pjreddie.com/darknet/yolo/,coco数据集下载
http://cocodataset.org/#overview,coco数据下载原址
https://blog.csdn.net/gqixf/article/details/79280224,coco数据集介绍
https://blog.csdn.net/hehangjiang/article/details/79108232#commentsedit(代码问题有点多)
http://www.cnblogs.com/arkenstone/p/7337077.html,将MSCOCO、Caltech转成VOC格式
https://blog.csdn.net/yjl9122/article/details/56842098,create xml代码参考