SSD系列(1)——制作自己的COCO格式Object Instance目标检测数据集

前言

目前在目标检测任务中，VOC和COCO是两个影响力最大的数据集系列。有时候我们希望将一些经典模型如RCNN、YOLO、SSD等运用到自己的数据集中，同时希望能够尽可能快的部署、验证，而不是花时间重新写一整套dataset和dataloader。此时我们可以将自己的数据集封装成VOC或COCO固定格式，然后对模型的接口做一些细微修改，就可以跑起来。
本文主要讨论COCO数据集的detection任务，并基于SSD-Pytorch的源码进行修改，将自己的数据集封装成COCO格式，并在SSD上跑起来。
关键词：COCO数据集、SSD、pytorch、目标检测

注：编者水平有限，如有谬误，欢迎指正。若要转载，请注明出处，谢谢。
联系方式：
邮箱：[email protected]
QQ：1156356625

---

COCO-detection数据集格式

注：格式相对官方做了简化处理
数据集的格式一般用一个json文件表示，训练集测试集分开保存。数据集格式如下：

generated_list = {"images": [],
                  "annotations": [],
                  "categories": []

拆开来，先说说images，它是一个字典的列表，其中每个元素对应一个样本图片:

{"file_name": 图片路径,
"width": 图片宽度,
"height": 图片高度,
"id": 图片编号}

前三项包含了文件的基本信息，id是最重要的一项，其编号是人为指定，与文件层级没有关系。
然后是annotations，同样也是一个字典的列表，每个元素对应一个标注。

{"image_id": 对应图片的编号,
"category_id": label, 
"id": annotation编号, 
"bbox": bbox}

这里的image_id和images的id要对应起来，这两个id是必须一对多严格对应起来的，否则在调试的时候可能会出现loss能下降但是不收敛的情况(找错标注)。label是标注类别，K类的话对应就是0到K-1。annotations也有一个id，这里的表示其自身的编号。因为一张图片对应多个标注框，所以annotation也要重新编号。bbox格式默认为XYWH格式，即[x, y, w, h]代表标注框左上角点x、y坐标与其宽高，其数据格式是列表。
最后是category的格式，同样也是字典列表。不过对于一个数据集，这项是定长的，可以预先生成好。

{"id": 0,		#类别编号
"name": "类别名0"}
 ...
 "id": K-1,
 "name": "类别名1"}

对于普通的检测任务，这些数据就已经足够了，接下来是封装。

自定义数据集封装

对于数据的排列没有要求，只需要图片和标注的文件在一个目录就可以，至于目录的层叠嵌套没有什么影响。封装可以分为两个步骤：

• 划分数据集

数据集的划分可以根据个人需求来，现在假设数据集没有人为划分过，按比例随机进行划分。我按照8：2把训练集划成train和val两部分，然后把划分的文件路径

# 过滤有效样本对，图片和标注需要在同一目录下，且命名前缀一致，如果残缺则滤掉
def get_path_list(root_dir):
	# 图片及标注格式
	img_extensions = ['jpg', 'jpeg', 'bmp', 'png', 'gif']
	annotation_extensions = ['json']
	# 图片及标注的路径列表
	images_path_list = []
	annotations_path_list = []
	# 遍历所有子路径
	for root, dir, files in os.walk(root_dir):
		if files:
			files.sort()
			imglist_org = [fn for fn in files if any(fn.endswith(ext) for ext in img_extensions)]
			anlist_org = []
			delete_index = []
			for i, l in enumerate(imglist_org):
				if l + "_.json" in files:
					anlist_org.append(l + "_.json")
				elif l + ".json" in files:
					anlist_org.append(l + ".json")
				else:
					delete_index.append(i)
			for i in delete_index[::-1]:
				imglist_org.pop(i)
			images_path_list += [os.path.join(root, fn) for fn in imglist_org]
			annotations_path_list += [os.path.join(root, fn) for fn in anlist_org]
	return images_path_list, annotations_path_list
# 按比例划分数据集
def split_trainval():
	train_root_dir = "../../datasets/mainland/train"
	images_path_list, annotations_path_list = get_path_list(train_root_dir)
	num_files = len(images_path_list)
	shuffle_list = [i for i in range(num_files)]
	split_ratio = [0.8, 0.2]
	typename = ["train", "val"]
	# shuffle索引
	split_index = ([0, 0] + [split_ratio[i] * num_files for i in range(2)])
	split_index = [int(split_index[ii] + split_index[ii + 1]) for ii in range(3)]
	split_index[-1] = -1
	random.shuffle(shuffle_list)
	for type in range(2):
		with open("./config/" + typename[type] + "_split.json", "w") as f:
			json.dump({
				"image_files": [images_path_list[i] for i in shuffle_list[split_index[type]:split_index[type + 1]]],
				"annotations_files": [annotations_path_list[i] for i in
									  shuffle_list[split_index[type]:split_index[type + 1]]]
			}, f)
			f.close()

• 生成json文件

单核跑太慢了，稍微改了下，写成了多进程的处理方式。不过这也引入了个坑，每个进程得到的annotations数目不固定，所以需要跑完重新编个号。

# num_works：要开的进程数
def get_dataset_plates(type,num_works=16):
	if type not in ["train", "val", "test"]:
		print("Error Dataset Type!")
		return
	root = "./config"
	with open(os.path.join(root, type + "_split.json")) as f:
		split_file = json.load(f)
		images_path_list = split_file["image_files"]
		annotations_path_list = split_file["annotations_files"]
		f.close()
	generated_list = {"images": [],
					  "annotations": [],
					  "categories": [
						  {"id": 1,
						   "name": "单层车牌"},
						  {"id": 2,
						   "name": "双层车牌"}]
					  }
	nums = images_path_list.__len__()
	idx = list(range(nums))
	# 开进程池
	p = Pool(num_works)
	result = []
	# 很暴力，直接把文件列表索引等分一下，然后每个进程扔整个列表进去，索引不同
	for i in range(num_works):
		index = idx[int(i*nums/num_works):int((i+1)*nums/num_works)]
		result.append(p.apply_async(get_dicts, args=(index, images_path_list, annotations_path_list)))
	# 等所有进程都跑完
	p.close()
	p.join()
	# 取每个进程的数据
	for p in result:
		im, an = p.get()
		generated_list["images"] += im
		generated_list["annotations"] += an
	# 重新给annotation编号
	for idx, ann in enumerate(generated_list["annotations"]):
		ann['id'] = idx
	with open("NP20w_" + type + ".json", "w") as f:
		json.dump(generated_list, f)
# 根据索引封装数据集
def get_dicts(index_list, images_path_list, annotations_path_list):
	'''
	我的数据集是车牌数据集，存储路径很杂乱，编码方式也不一样，同时标注的json也有两种格式
	实际处理需要根据实际情况，这部分代码仅作参考
	'''
	images = []
	annotations =  []
	for index in index_list:
		file = images_path_list[index]
		annotations_file = annotations_path_list[index] if file in annotations_path_list[index] else print(
			"Invalid Annotation File.")
		# 两次打开文件，第一次是用chardet.detect检查编码格式，第二次是正式打开
		with open(annotations_file, mode="rb") as f:
			encoding = chardet.detect(f.read()).get("encoding")
			f.close()
		with open(annotations_file, encoding=encoding) as f:
			json_dict = json.load(f)
		# 两种标注格式检查一下
		if "plate_count" not in json_dict.keys():
			annotations_type = 0
			plate_count = json_dict["Plates"].__len__()
			plate_pos = ["PosPt0X", "PosPt0Y", "PosPt3X", "PosPt3Y"]
		else:
			annotations_type = 1
			plate_count = json_dict["plate_count"]
			plate_pos = ["Pt0X", "Pt0Y", "Pt2X", "Pt2Y"]
		for i in range(plate_count):
			plate_dict = json_dict["plate_" + str(i)] if annotations_type else json_dict["Plates"][i]
			plate_type = plate_dict["plateType"] if annotations_type else plate_dict["PlateType"]
			# 类别数只有两类
			if "双" in plate_type:
				plate_type = 1
			else:
				plate_type = 0
			# 生成车牌标注
			bbox = []
			for pt_key in plate_pos:
				if annotations_type:
					bbox.append(plate_dict[pt_key] / 10)
				else:
					bbox.append(float(plate_dict[pt_key]))
			# 很关键，bbox要好好检查，不然出现错误框，训练时会随机在某个iter梯度nan掉
			if bbox[0] >= bbox[2] or bbox[1] >= bbox[3]:
				print("Error bbox!")
				continue
			# 这里将XYXY格式转成XYWH格式，用了detectron2的库，也可以简单的加减处理一下
			bbox = boxes.BoxMode.convert(bbox, boxes.BoxMode.XYXY_ABS, boxes.BoxMode.XYWH_ABS)
			annotations.append(
				{"image_id": index, "category_id": plate_type, "id": 0, "bbox": bbox})
		im = Image.open(file)
		images.append({
			"file_name": file,
			"width": im.size[0],
			"height": im.size[1],
			"id": index
		})
	return images, annotations

好啦，数据集处理完啦。其实做实际项目的时候，大部分时间都是在跟数据杠。数据残缺、错误、污染可能会导致训练的时候出现各种问题，所以数据源需要认真清洗，遇到问题也不妨多观察一下数据集。
COCO数据集的封装就讲到这里，下一篇是SSD-Pytorch调用数据集的API，请多多指教！