CenterNet 训练自己的数据
- 1.环境配置
- 2.数据准备
- 3. 修改配置
- 测试
1.环境配置
作者在github 上给出的安装步骤已经较为详细，可以参照作者的步骤进行安装。
code:https://github.com/xingyizhou/CenterNet
在这个过程中比较麻烦的是，作者的环境为pytorch 0.4.1 cuda9.0 cudnn7.1.2，cuda版本不匹配的话，后面会报错：ImportError:/home/shiep/CenterNet/src/lib/models/networks/DCNs/_ext/dcn_v2/dcn_v2.so: undefined symbol: __cudaRegisterFatBinaryEnd。
我使用的环境是 torch1.2 cuda10.0 在进行到作者的第5步Compile deformable convolutional (from DCNv2).时，将DCNv2的文件夹替换，可以去官网下载，也可以直接git clone https://github.com/CharlesShang/DCNv2.git,然后./make.sh后续操作和作者一致。
2.数据准备
由于训练使用coco数据的格式，所以要先将自己的数据格式转为coco的形式，我的数据形式是yolo的txt形式,txt存放的是：类别，x, y, w , h。通过下面的脚本将txt转成：图片地址，类别，xmin,ymin，xmax ,ymax的txt。
脚本如下：
import os
import cv2
# 标签路径，yolo的txt
originLabelsDir = r'/data_1/project/pingxiang/models_data/darknet/data/names/names/val_label'
#保存txt路径
saveDir = r'/data_1/project/pingxiang/models_data/darknet/data/names/names/txt/val.txt'
# 图片路径
originImagesDir = r'/data_1/project/pingxiang/models_data/darknet/data/names/names/val_img'
txtFileList = os.listdir(originLabelsDir)
with open(saveDir, 'w') as fw:
for txtFile in txtFileList:
with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:
labelList = fr.readlines()
for label in labelList:
label = label.strip().split()
x = float(label[1])
y = float(label[2])
w = float(label[3])
h = float(label[4])
imagePath = os.path.join(originImagesDir,txtFile.replace('txt', 'jpg'))
image = cv2.imread(imagePath)
H, W, _ = image.shape
x1 = (x - w / 2) * W
y1 = (y - h / 2) * H
x2 = (x + w / 2) * W
y2 = (y + h / 2) * H
fw.write(txtFile.replace('txt', 'jpg') + ' {} {} {} {} {}\n'.format(int(label[0]) + 1, x1, y1, x2, y2))
print('{} done'.format(txtFile))将得到上述的txt后，再转为coco的.json文件。脚本如下：
import json
import os
import cv2
#root_path包含images(图片文件夹)，val.txt(bbox标注,上述生成的txt)，classes.txt(类别标签,一行一个类别，例如：car people),最终的结果保存在annotations文件夹下。
root_path = r'/data_1/project/pingxiang/models_data/darknet/data/names/names'
phase = 'val'
dataset = {'categories': [], 'annotations': [], 'images': []}
with open(os.path.join(root_path, 'classes.txt')) as f:
classes = f.read().strip().split()
for i, cls in enumerate(classes, 1):
dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'})
indexes = os.listdir(os.path.join(root_path, 'images'))
global indx
indx= 0
with open(os.path.join(root_path, 'val.txt')) as tr:
annos = tr.readlines()
for k, index in enumerate(indexes):
count += 1
im = cv2.imread(os.path.join(root_path, 'images/') + index)
height, width, _ = im.shape
dataset['images'].append({'file_name': index,
'id': k,
'width': width,
'height': height})
for ii, anno in enumerate(annos):
parts = anno.strip().split()
if parts[0] == index:
cls_id = parts[1]
x1 = float(parts[2])
y1 = float(parts[3])
x2 = float(parts[4])
y2 = float(parts[5])
width = max(0, x2 - x1)
height = max(0, y2 - y1)
dataset['annotations'].append({
'area': width * height,
'bbox': [x1, y1, width, height],
'category_id': int(cls_id),
'id': i,
'image_id': k,
'iscrowd': 0,
'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
})
print('{} images handled'.format(indx))
folder = os.path.join(root_path, 'annotations')
if not os.path.exists(folder):
os.makedirs(folder)
json_name = os.path.join(root_path, 'annotations/{}.json'.format(phase))
with open(json_name, 'w') as f:
json.dump(dataset, f)需要注意的是，手动划分一下train test val的图片标签, 各自生成。最后得到三个文件：train.json test.json val.json。得到这些文件后，放在CenterNet/data/下，annotations存放json文件，images存放图片。
3. 修改配置
修改src/lib/datasets/dataset文件下的coco.py：
1）num_classes=80改成自己的类别数
2）修改数据和图片路径，data_dir 自己数据集文件夹的名字，img_dir 是 images 图片文件夹
3）类别名字和类别id改成自己对应的类别名和ID
修改src/lib/utils/debugger.py
1）修改自己的类别名和类别数
在正式训练之前建议debug到coco.p如下图这部分，看一看读取train和val部分是不是正常读取。
def __init__(self, opt, split):
super(COCO, self).__init__()
self.data_dir = os.path.join(opt.data_dir, 'hail')
self.img_dir = os.path.join(self.data_dir, 'images')
if split == 'val':
self.annot_path = os.path.join(
self.data_dir, 'annotations',
'test.json').format(split)
else:
if opt.task == 'ctdet':
self.annot_path = os.path.join(
self.data_dir, 'annotations',
'train.json').format(split)
else:
self.annot_path = os.path.join(
self.data_dir, 'annotations',
'train.json').format(split)一切正常后，训练指令：
python main.py ctdet --exp_id coco_dla --batch_size 16 --master_batch 1 --lr 1.25e-4 --gpus 0测试
运行demo.py ，测试时注意将类别数和对应的ID改成自己需要的。
更新:
xml to coco 脚本
# coding:utf-8
import xml.etree.ElementTree as ET
import os
import json
coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []
category_set = dict()
image_set = set()
category_item_id = -1
image_id = 20180000000
annotation_id = 0
def addCatItem(name):
global category_item_id
category_item = dict()
category_item['supercategory'] = 'none'
# category_item_id += 1
# '''
# 11 label
# '''
# if name == "smallCar":
# category_item_id = 0 #四轮车
# if name == "bigCar":
# category_item_id = 1 #罚单
# if name == "people":
# category_item_id = 2 #交警
# if name == "xingren-qiche":#二轮车
# category_item_id = 3
# if name =="feijidongche":#人脸
# category_item_id = 4
# if name=="xingren":#普通人
# category_item_id=5
# if name=="xingren-tuiche":#摔倒的人
# category_item_id=6
# if name=="feijidongche_1":#测酒仪
# category_item_id=7
# if name == "car_other":#三轮车
# category_item_id = 8
# if name == "car_stop":#签字
# category_item_id = 9
# if name == "car_direct_other":#驾驶证
# category_item_id = 10
# '''
# 7类
# '''
# if name == "bigCar":
# category_item_id = 0
# if name=="feijidongche_1":
# category_item_id=1
# if name == "car_stop":
# category_item_id = 2
# if name =="feijidongche":
# category_item_id = 3
# if name == "car_front":#车窗
# category_item_id = 4
# if name == "car_back":#打票机
# category_item_id = 5
# if name == "car_left":#车窗上的罚单
# category_item_id = 6
if name == "Car":
category_item_id = 0
if name=="people":
category_item_id=1
if name =="machine":
category_item_id = 2
if name =="paper":
category_item_id = 3
if name == "car_stop":#签字
category_item_id = 4
category_item['id'] = category_item_id
category_item['name'] = name
coco['categories'].append(category_item)
category_set[name] = category_item_id
return category_item_id
def addImgItem(file_name, size):
global image_id
if file_name is None:
raise Exception('Could not find filename tag in xml file.')
if size['width'] is None:
raise Exception('Could not find width tag in xml file.')
if size['height'] is None:
raise Exception('Could not find height tag in xml file.')
image_id += 1
image_item = dict()
image_item['id'] = image_id
image_item['file_name'] = file_name
image_item['width'] = size['width']
image_item['height'] = size['height']
coco['images'].append(image_item)
image_set.add(file_name)
return image_id
def addAnnoItem(object_name, image_id, category_id, bbox):
global annotation_id
annotation_item = dict()
annotation_item['segmentation'] = []
seg = []
# bbox[] is x,y,w,h
# left_top
seg.append(bbox[0])
seg.append(bbox[1])
# left_bottom
seg.append(bbox[0])
seg.append(bbox[1] + bbox[3])
# right_bottom
seg.append(bbox[0] + bbox[2])
seg.append(bbox[1] + bbox[3])
# right_top
seg.append(bbox[0] + bbox[2])
seg.append(bbox[1])
annotation_item['segmentation'].append(seg)
annotation_item['area'] = bbox[2] * bbox[3]
annotation_item['iscrowd'] = 0
annotation_item['ignore'] = 0
annotation_item['image_id'] = image_id
annotation_item['bbox'] = bbox
annotation_item['category_id'] = category_id
annotation_id += 1
annotation_item['id'] = annotation_id
coco['annotations'].append(annotation_item)
def parseXmlFiles(xml_path):
for f in os.listdir(xml_path):
if not f.endswith('.xml'):
continue
bndbox = dict()
size = dict()
current_image_id = None
current_category_id = None
file_name = None
size['width'] = None
size['height'] = None
size['depth'] = None
xml_file = os.path.join(xml_path, f)
print(xml_file)
tree = ET.parse(xml_file)
root = tree.getroot()
if root.tag != 'annotation':
raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
# elem is <folder>, <filename>, <size>, <object>
for elem in root:
current_parent = elem.tag
current_sub = None
object_name = None
if elem.tag == 'folder':
continue
if elem.tag == 'filename':
file_name = elem.text
if file_name in category_set:
raise Exception('file_name duplicated')
# add img item only after parse <size> tag
elif current_image_id is None and file_name is not None and size['width'] is not None:
if file_name not in image_set:
current_image_id = addImgItem(file_name, size)
print('add image with {} and {}'.format(file_name, size))
else:
raise Exception('duplicated image: {}'.format(file_name))
# subelem is <width>, <height>, <depth>, <name>, <bndbox>
for subelem in elem:
bndbox['xmin'] = None
bndbox['xmax'] = None
bndbox['ymin'] = None
bndbox['ymax'] = None
current_sub = subelem.tag
if current_parent == 'object' and subelem.tag == 'name':
object_name = subelem.text
if object_name not in category_set:
current_category_id = addCatItem(object_name)
else:
current_category_id = category_set[object_name]
elif current_parent == 'size':
if size[subelem.tag] is not None:
raise Exception('xml structure broken at size tag.')
size[subelem.tag] = int(subelem.text)
# option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
for option in subelem:
if current_sub == 'bndbox':
if bndbox[option.tag] is not None:
raise Exception('xml structure corrupted at bndbox tag.')
bndbox[option.tag] = int(option.text)
# only after parse the <object> tag
if bndbox['xmin'] is not None:
if object_name is None:
raise Exception('xml structure broken at bndbox tag')
if current_image_id is None:
raise Exception('xml structure broken at bndbox tag')
if current_category_id is None:
raise Exception('xml structure broken at bndbox tag')
bbox = []
# x
bbox.append(bndbox['xmin'])
# y
bbox.append(bndbox['ymin'])
# w
bbox.append(bndbox['xmax'] - bndbox['xmin'])
# h
bbox.append(bndbox['ymax'] - bndbox['ymin'])
print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
bbox))
addAnnoItem(object_name, current_image_id, current_category_id, bbox)
if __name__ == '__main__':
xml_path = '/data_1/project/nanchang/Dataset/五大类/1-3批数据/centernet/xml_5类/val/'
json_file = '/data_1/project/nanchang/Dataset/五大类/1-3批数据/centernet/val.json'
parseXmlFiles(xml_path)
json.dump(coco, open(json_file, 'w'))
本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。
