之前的DAGM2COCO代码有一点问题,重新上传了新的代码
前言
因为在做目标检测时需要用到DAGM2007数据集,但是官方提供的数据集格式与现在主流的COCO、VOC格式不同,重新写Dataset类又太过麻烦。因此在这里写了几个代码方便进行数据集格式转换。
以下代码如果有不规范或错误的地方欢迎大家批评指正!
注意:本博客中的COCO2VOC以及VOC2COCO仅针对灰度图使用,若需要转换COCO数据集或Pascal VOC数据集的格式,需要对代码进行修改。
另外参考了这位大佬写的文章,写的真的很清楚
VOC格式与COCO格式详解
1. DAGM2007数据集下载
DAGM2007数据集包含10个文件夹,前6个是训练集,后4个是测试集,常被用于深度学习的目标检测中评估网络的性能。下面是DAGM2007的下载链接,填写邮箱信息后会自动发送下载链接到该邮箱。
DAGM2007下载地址
有时候可能进不去,多刷两遍就可以进去了。实在进不去可以看看其他博客,有大佬上传到百度网盘上了。
2. DAGM2COCO
逻辑大概是:
1、先从原数据集中读取图像所需要的信息
2、将这些信息写成COCO的格式,写入新数据集中
3、复制图像到新数据集
最后新路径下数据集的格式为:
- DAGM-COCO
- annotations
- instances_train.json
- instances_val.json
- train
- images
- val
- images
- annotations
代码
这里因为DAGM2007数据集中图像的名字重复,因此重新命了个名
obj_list是十个类名
import os
import cv2
import json
import shutil
obj_list = ['class1', 'class2', 'class3', 'class4', 'class5', 'class6', 'class7', 'class8', 'class9', 'class10']
# 读取数据集的图片信息(图片路径、种类、标签路径,另外为每张图像编了个号)
# 另外将图片复制到新的文件夹
def get_infor(path, mode, new_dataset_path, rate=0.5):
# 获取数据集内所有图片信息
information = []
dataset_cla = [cla for cla in os.listdir(path) if os.path.isdir(os.path.join(path, cla))]
i = 0
for cla in dataset_cla:
cla_path = os.path.join(path, cla, cla)
f = open(os.path.join(cla_path, mode, 'Label', 'Labels.txt'), 'r')
f_train = f.read().split('\n')
for line in f_train:
img = {}
if len(line) <= 1:
continue
infor = line.split('\t')
if int(infor[1]) == 1:
img['image_path'] = os.path.join(cla_path, mode, infor[0] + '.PNG')
img['category_id'] = int(cla.strip('Class'))
img['label_path'] = os.path.join(cla_path, mode, 'Label', infor[4])
img['image_id'] = i
information.append(img)
i += 1
f.close()
# 这一部分是可以优化的,但是为了防止出现错误,在这里又重新加了个for循环
if mode == 'Train':
for img in information:
new_name = ''.join('0' for i in range(8 - len(str(img['image_id'])))) + str(img['image_id']) + '.PNG'
img['file_name'] = new_name
shutil.copy(img['image_path'], os.path.join(new_dataset_path, 'train', new_name))
return information
else:
for img in information:
new_name = ''.join('0' for i in range(8 - len(str(img['image_id'])))) + str(img['image_id']) + '.PNG'
img['file_name'] = new_name
shutil.copy(img['image_path'], os.path.join(new_dataset_path, 'val', new_name))
return information
# 获取bbox的四个值
def get_bbox(label_path):
annotations = []
img = cv2.imread(label_path)[:, :, 0]
contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
# 每条轮廓挨个处理
# 因为我的图像只有一条轮廓,所以没有这个循环也可以
for cnt in contours:
# 得到外接矩形(正的)
x, y, w, h = cv2.boundingRect(cnt) # xy为左上角坐标,wh为宽高;COCO读出来的也是这个格式
annotations.append([x, y, w, h])
return annotations
# 将读取到的图片信息处理,写成COCO数据集的格式
def infor2coco(infor, new_path, mode):
coco = {}
info = {}
images = []
annotations = []
categories = []
global_annotation_id = 0
info['year'] = 2007
info['contributor'] = 'https://hci.iwr.uni-heidelberg.de/content/weakly-supervised-learning-industrial-optical-inspection'
coco['info'] = info
for img_infor in infor:
image = {'id': img_infor['image_id']}
img = cv2.imread(img_infor['image_path'])[:, :, 0]
image['width'] = img.shape[0]
image['height'] = img.shape[1]
image['file_name'] = img_infor['file_name']
images.append(image)
bboxes = get_bbox(img_infor['label_path'])
if len(bboxes) > 1:
for bbox in bboxes:
sub_annotation = {'id': global_annotation_id,
'image_id': img_infor['image_id'],
'category_id': img_infor['category_id'],
'bbox': bbox, 'iscrowd': 0,
'area': bbox[2] * bbox[3]}
annotations.append(sub_annotation)
global_annotation_id += 1
else:
annotation = {'id': global_annotation_id,
'image_id': img_infor['image_id'],
'category_id': img_infor['category_id'],
'bbox': bboxes[0], 'iscrowd': 0,
'area': bboxes[0][2] * bboxes[0][3]}
annotations.append(annotation)
global_annotation_id += 1
coco['images'] = images
coco['annotations'] = annotations
for i, cla in enumerate(obj_list):
category = {'id': i + 1,
'name': cla,
'supercategory': 'defect'}
categories.append(category)
coco['categories'] = categories
file_name = f'{new_path}/annotations/instances_{mode}.json'
if os.path.exists(file_name):
os.remove(file_name)
json.dump(coco, open(file_name, 'w'))
def main():
path = 'dataset/DAGM2007'
new_path = 'dataset/DAGM2007-coco'
infor = get_infor(path, 'Train', new_dataset_path=new_path)
infor2coco(infor, new_path, 'train')
val_infor = get_infor(path, 'Test', new_dataset_path=new_path)
infor2coco(val_infor, new_path, 'val')
if __name__ == '__main__':
main()
3. DAGM2VOC
有大佬写了,这里放上大佬的链接
DAGM2VOC
4. COCO2VOC
这里参考了另外一位博主的方法,下面是他的链接
COCO2VOC
代码
from pycocotools.coco import COCO
import os
import shutil
from lxml import etree, objectify
from tqdm import tqdm
# 写入xml格式的标注信息
def write_xml(root, annotation, ori_root, t):
filename = annotation["filename"] if t == 'train' else str(int(annotation["filename"].split('.')[0]) + 10000000) + '.PNG'
annot_save_path = os.path.join(root, 'Annotations', filename.replace('PNG', 'xml'))
img_save_path = os.path.join(root, 'JPEGImages', filename)
shutil.copy(os.path.join(ori_root, t, annotation["filename"]),
img_save_path)
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder(root),
E.filename(filename),
E.size(
E.width(annotation['size'][0]),
E.height(annotation['size'][1]),
E.depth(annotation['size'][2])
),
E.segmented(0)
)
for obj in annotation['objects']:
E2 = objectify.ElementMaker(annotate=False)
anno_tree2 = E2.object(
E.name(obj['name']),
E.pose('Unspecified'),
E.truncated(0),
E.difficult(0),
E.bndbox(
E.xmin(obj['bbox'][0]),
E.ymin(obj['bbox'][1]),
E.xmax(obj['bbox'][2]),
E.ymax(obj['bbox'][3])
)
)
anno_tree.append(anno_tree2)
etree.ElementTree(anno_tree).write(annot_save_path, pretty_print=True)
return filename
def trans_c2v(coco, index, classes, voc_root, coco_root, t):
img_info = coco.loadImgs(index)[0]
annots = coco.loadAnns(coco.getAnnIds(imgIds=index))
annotation = dict()
annotation['filename'] = img_info['file_name']
annotation['size'] = [img_info['width'], img_info['height'], 3]
objects = []
for annot in annots:
obj = dict()
obj['name'] = classes[annot['category_id']]
obj['bbox'] = [annot['bbox'][0],
annot['bbox'][1],
annot['bbox'][0] + annot['bbox'][2],
annot['bbox'][1] + annot['bbox'][3]]
objects.append(obj)
annotation['objects'] = objects
filename = write_xml(voc_root, annotation, coco_root, t)
return filename
def main(coco_root, voc_root):
voc_imgs_path = os.path.join(voc_root, 'JPEGImages')
voc_annots_path = os.path.join(voc_root, 'Annotations')
voc_set_path = os.path.join(voc_root, 'ImageSets', 'Main')
if not os.path.exists(voc_root):
os.makedirs(voc_imgs_path)
os.makedirs(voc_annots_path)
os.makedirs(voc_set_path)
else:
shutil.rmtree(voc_root)
main(coco_root, voc_root)
train_txt = []
val_txt = []
train_coco = COCO(os.path.join(coco_root, 'annotations', 'instances_train.json'))
category = train_coco.dataset['categories']
classes = dict()
for c in category:
classes[c['id']] = c['name']
images_ID = train_coco.getImgIds()
for ids in tqdm(images_ID):
filename = trans_c2v(train_coco, ids, classes, voc_root, coco_root, 'train')
train_txt.append(filename)
val_coco = COCO(os.path.join(coco_root, 'annotations', 'instances_val.json'))
images_ID = val_coco.getImgIds()
for ids in tqdm(images_ID):
filename = trans_c2v(val_coco, ids, classes, voc_root, coco_root, 'val')
val_txt.append(filename)
with open(os.path.join(voc_set_path, 'train.txt'), 'w') as f:
for i, name in enumerate(train_txt):
f.write(name.split('.')[0])
if i <= (len(train_txt) - 1):
f.write('\n')
with open(os.path.join(voc_set_path, 'val.txt'), 'w') as f:
for i, name in enumerate(val_txt):
f.write(name.split('.')[0])
if i <= (len(val_txt) - 1):
f.write('\n')
if __name__ == '__main__':
# COCO格式数据集存放地址
coco_root = '/DAGM2007-COCO'
# VOC格式数据集地址
voc_root = '/DAGM2007-VOC'
main(coco_root, voc_root)
测试代码
from pycocotools.coco import COCO
import os
import cv2
root = '/DAGM2007-COCO'
set_name = 'train'
dagm_coco = COCO(os.path.join(root, 'annotations', 'instances_' + set_name + '.json'))
image_ids = dagm_coco.getImgIds()
for index in image_ids:
image_info = dagm_coco.loadImgs(index)[0]
annots = dagm_coco.loadAnns(dagm_coco.getAnnIds(imgIds=index))
image_path = os.path.join(root, set_name, image_info['file_name'])
img = cv2.imread(image_path)
for annot in annots:
cv2.rectangle(img,
(annot['bbox'][0], annot['bbox'][1]),
(annot['bbox'][0] + annot['bbox'][2], annot['bbox'][1] + annot['bbox'][3]),
(255, 0, 0), 1)
cv2.imshow(f"{image_info['file_name']}", img)
cv2.waitKey(1000)
cv2.destroyAllWindows()