1、下載COCO數據集：
https://blog.csdn.net/aLWX_hust/article/details/88350248?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522158541005019726867813769%2522%252C%2522scm%2522%253A%252220140713.130056874…%2522%257D&request_id=158541005019726867813769&biz_id=0&utm_source=distribute.pc_search_result.none-task
下載你需要的，博主下載的數據集類如下：
一共三份，用迅雷下！！！
2、解壓：
解壓到同個文件夾下：
3、 提取：
from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw
import random
headstr = """\
<annotation>
<folder>VOC</folder>
<filename>%s</filename>
<source>
<database>My Database</database>
<annotation>COCO</annotation>
<image>flickr</image>
<flickrid>NULL</flickrid>
</source>
<owner>
<flickrid>NULL</flickrid>
<name>company</name>
</owner>
<size>
<width>%d</width>
<height>%d</height>
<depth>%d</depth>
</size>
<segmented>0</segmented>
"""
objstr = """\
<object>
<name>%s</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>%d</xmin>
<ymin>%d</ymin>
<xmax>%d</xmax>
<ymax>%d</ymax>
</bndbox>
</object>
"""
tailstr = '''\
</annotation>
'''
# 如果目錄不存在則創建它，存在則刪除後創建
def mkr(path):
if os.path.exists(path):
# 刪除後創建
shutil.rmtree(path)
# os.mkdir(path)
# os.mkdir創建單層目錄；os.makedirs創建多層目錄
os.makedirs(path)
else:
os.makedirs(path)
# 通過coco數據集的id，得到它的類別名
def id2name(coco):
classes = dict()
for cls in coco.dataset['categories']:
classes[cls['id']] = cls['name']
return classes
# 寫xml文件
def write_xml(anno_path, head, objs, tail):
f = open(anno_path, "w")
f.write(head)
for obj in objs:
f.write(objstr % (obj[0], obj[1], obj[2], obj[3], obj[4]))
f.write(tail)
#
def save_annotations_and_imgs(coco, dataset, filename, objs):
# eg:COCO_train2014_000000196610.jpg-->COCO_train2014_000000196610.xml
anno_path = anno_dir + filename[:-3] + 'xml'
print(anno_path)
img_path=dataDir+dataset+'/'+filename
#img_path = dataset + '//' + filename
print(img_path)
dst_imgpath = img_dir + filename
# print(dst_imgpath)
img = cv2.imread(img_path)
# print("***********")
# print(img1.shape)
try:
if (img.shape[2] == 1):
print(filename + " not a RGB image")
return
except:
print(img_path)
shutil.copy(img_path, dst_imgpath)
head = headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
tail = tailstr
write_xml(anno_path, head, objs, tail)
def showimg(coco, dataset, img, classes, cls_id, show=True):
global dataDir
I = Image.open('%s/%s/%s' % (dataDir, dataset, img['file_name']))
# print(I)#
# 通過id，得到註釋的信息
annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
# print(annIds)
anns = coco.loadAnns(annIds)
# print(anns)
# coco.showAnns(anns)
objs = []
for ann in anns:
class_name = classes[ann['category_id']]
# print(class_name)
if class_name in classes_names:
# print(class_name)
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 = [class_name, xmin, ymin, xmax, ymax]
objs.append(obj)
draw = ImageDraw.Draw(I)
draw.rectangle([xmin, ymin, xmax, ymax])
if show:
plt.figure()
plt.axis('off')
plt.imshow(I)
plt.show()
return objs
def generate():
for dataset in datasets_list:
# ./COCO/annotations/instances_train2014.json
annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataset)
# COCO API for initializing annotated data
coco = COCO(annFile)
'''
COCO 對象創建完畢後會輸出如下信息:
loading annotations into memory...
Done (t=0.81s)
creating index...
index created!
至此, json 腳本解析完畢, 並且將圖片和對應的標註數據關聯起來.
'''
# show all classes in coco
classes = id2name(coco)
print(classes)
# [1, 2, 3, 4, 6, 8]
classes_ids = coco.getCatIds(catNms=classes_names)
print(classes_ids)
# print("**********")
print(classes_names)
for cls in classes_names:
# Get ID number of this class
cls_id = coco.getCatIds(catNms=[cls])
img_ids = coco.getImgIds(catIds=cls_id)
# print(cls,len(img_ids))
# imgIds=img_ids[0:10]
for imgId in tqdm(img_ids):
img = coco.loadImgs(imgId)[0]
# print("##########")
# print(img)
# print("##########")
filename = img['file_name']
# print(filename)
objs = showimg(coco, dataset, img, classes, classes_ids, show=False)
# print(objs)
save_annotations_and_imgs(coco, dataset, filename, objs)
def split_traintest(trainratio=0.7, valratio=0.2, testratio=0.1):
dataset_dir = dataDir
files = os.listdir(img_dir)
trains = []
vals = []
trainvals = []
tests = []
random.shuffle(files)
for i in range(len(files)):
filepath = img_dir + "/" + files[i][:-3] + "jpg" # 找到以上images文件夾下的圖片
# print(filepath)
if (i < trainratio * len(files)):
trains.append(files[i])
trainvals.append(files[i])
elif i < (trainratio + valratio) * len(files):
vals.append(files[i])
trainvals.append(files[i])
else:
tests.append(files[i])
# 生成yolo所用的txt
with open(dataset_dir + "/trainval.txt", "w")as f:
for line in trainvals:
line = img_dir + "/" + line
f.write(line + "\n")
with open(dataset_dir + "/test.txt", "w") as f:
for line in tests:
line = img_dir + "/" + line
f.write(line + "\n")
# 生成voc所用的txt
maindir = dataset_dir + "ImageSets/Main"
mkr(maindir)
with open(maindir + "/train.txt", "w") as f:
for line in trains:
line = line[:line.rfind(".")]
f.write(line + "\n")
with open(maindir + "/val.txt", "w") as f:
for line in vals:
line = line[:line.rfind(".")]
f.write(line + "\n")
with open(maindir + "/trainval.txt", "w") as f:
for line in trainvals:
line = line[:line.rfind(".")]
f.write(line + "\n")
with open(maindir + "/test.txt", "w") as f:
for line in tests:
line = line[:line.rfind(".")]
f.write(line + "\n")
print("spliting done")
if __name__ == "__main__":
# 下載的coco數據集存放路徑，annotations和train2014/val2014/...
dataDir = 'F:/coco/'
# 轉化後的文件保存路徑
savepath = "F:/sheep/"
# 轉化後圖片保存路徑
img_dir = savepath + 'images/'
# 轉化後xml保存路徑
anno_dir = savepath + 'Annotations/'
# 選用的數據集列表
# datasets_list=['train2014', 'val2014']
datasets_list = ['train2017']
# 想要從coco數據集中提取的類別
# coco數據集目標檢測中有90個id，但實際只有80類
classes_names = ['sheep']
"""classes_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag',
'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table',
'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush']"""
# 創建所需要的目錄
mkr(savepath)
mkr(img_dir)
mkr(anno_dir)
# 生成voc格式的數據集
generate()
# 生成yolo及voc所使用的txt文件
split_traintest()結果如下：
溫馨提示：博主提取後訓練數據，發現總是出現bug，查到最後是關於數據集中包含黑白照片的問題，雖然影響可能不會很大，但是還是建議刪除黑白照片，留下RGB格式的，以下是黑白照片的數字後幾位：
黑白：01645，35732，43773，46382，48421，48907，49458，78511，85657，120666124694，146202，143382，167369，193951，202772，204792，206473，240054，103008260990，261068，269532，268741，275829，277854，280999，315668，314065，303438312876，319163，321328，327834，347179，394208，421195，428445，432917，27883445952，456439，468085，477639，477679，481098，480588，486641，496636496682，502629，516319，535010，537499，548951，553488，427310，409808