centernet-训练和验证自己的数据集(安全帽)
本文以安全帽检测为例，验证下训练自己的数据集
1. 创建json数据
将自己的coco类型的数据放到根目录的dataset目录下
~/AI/CenterNet-master/data/helmet$ ls
annotations images kmeans.py meanvalue.txt
/annotations$ ls
instances_test2019.json instances_train2019.json instances_val2019.json
-
annotations的json文件，文件内容描述了各个数据集的数据相关信息内容，包括：标注信息，图片文件名称，图片路径等
-
images：所有的图片文件，代码解析json后，回到该目录查找图片相关内容
-
kmeans.py计算所有图片的均值和标准值文件
2. 生成图片的均值和标准值
//kmeans.py
import cv2, os, argparse
import numpy as np
from tqdm import tqdm
def main():
dirs = 'images'
img_file_names = os.listdir(dirs)
m_list, s_list = [], []
for img_filename in tqdm(img_file_names):
img = cv2.imread(dirs + '/' + img_filename)
//对读取的图片数据，归一化到[0,1]
img = img / 255.0
//计算均值和均方差，每个通道一个值
m, s = cv2.meanStdDev(img)
m_list.append(m.reshape((3,)))
s_list.append(s.reshape((3,)))
m_array = np.array(m_list)
s_array = np.array(s_list)
//求均值
m = m_array.mean(axis=0, keepdims=True)
s = s_array.mean(axis=0, keepdims=True)
print("mean = ", m[0][::-1])
print("std = ", s[0][::-1])
if __name__ == '__main__':
main()
~/AI/CenterNet-master/data/helmet$ python kmeans.py
mean = [0.5684546 0.53629323 0.49678354]
std = [0.25609233 0.25254378 0.26302197]
3. 创建和编辑数据加载类
拷贝coco.py为一个helmet.py名称，helmet最好以你的场景命名
~/AI/CenterNet-master/src/lib/datasets/dataset$ cp coco.py helmet.py
主要修改一下地方：
1.修改class名称为Helmet
2.修改num_classes为2，因为我们是二分类检测
3.修改数据集的均值和均方差为你自己的数据集的值
4.修改super类的参数为Helmet
5.修改数据集的目录的图片目录
6.修改测试集，验证集和训练集的json文件加载目录
7.修改分类名称为你自己的[‘nothat’, ‘wearhat’]
8.修改分类名称对应的id，要从0开始，0代表背景分类（一定要包含背景，网上其他参考文件从1开始，训练的时候分类名称映射会报错）
在dataset_factory中加入你的数据集
修改位置将数据集加入src/lib/datasets/dataset_factory.py
主要修改一下两个地方
+ from .dataset.helmet import Helmet
dataset_factory = {
'coco': COCO,
'pascal': PascalVOC,
'kitti': KITTI,
'coco_hp': COCOHP ,
+ 'helmet': Helmet
}
/src/lib/opts.py
修改模式使用的数据库，不修改也没有问题，训练的时候指定就可以了
self.parser.add_argument('--dataset', default='helmet',
help='coco | kitti | coco_hp | pascal | helmet')
修改默认的数据集信息,主要是input_shape大小，分类格式，数据集均值和标准值
def init(self, args=''):
337 default_dataset_info = {
338 'ctdet': {'default_resolution': [512, 512], 'num_classes': 2,
339 'mean': [0.568, 0.536, 0.496], 'std': [0.256, 0.252, 0.263],
340 'dataset': 'helmet'},
src/lib/utils/debugger.py
添加以下代码，主要是将数据集名称与分类名称进行绑定
elif dataset == 'helmet':
self.names = helmet_class_name
helmet_class_name = ['nothat','wearhat']
4.训练数据
#从随机参数开始训练
python main.py ctdet --exp_id helmet --batch_size 32 --lr 1.25e-4 --gpus 0,1 --num_workers 4 --arch res_18 --dataset helmet --input_h 608 --input_w 608 --lr_step 90,120 --num_epochs 200
#如果在某个模型的基础上预训练
--load_model ../models/ctdet_dla_2x.pth
- lr 学习率
- gpus 0,1 使用的gpu显卡
- num_workers 4 num_worker设置得大，好处是寻batch速度快
-batch_size 批量处理的图片的大小默认为32 - arch res_18 算法要使用的backbone
- dataset helmet 要训练的数据集
- input_h --input_w
- lr_step 90,120 学习率变化扩大，没迭代到一个lr_step值，学习率下降到原来的0.1
- num_epochs 200 总共的训练的epchs的次数。每将全部数据训练一次，表示一个epochs;注意与iterations区分开，iterations表示，每将batch size张图片训练一次，表示一个iteration
成功运行的日志为：
python main.py ctdet --exp_id helmet --batch_size 32 --lr 1.25e-4 --gpus 0,1 --num_workers 4 --arch res_18 --dataset helmet --input_h 608 --input_w 608 --lr_step 90,120 --num_epochs 200
certernet) os@os-Super-Server:~/AI/CenterNet-master/src$ python main.py ctdet --exp_id helmet --batch_size 32 --lr 1.25e-4 --gpus 0,1 --num_workers 4 --arch res_18 --dataset helmet --input_h 608 --input_w 608 --lr_step 90,120 --num_epochs 200
Fix size testing.
training chunk_sizes: [16, 16]
#模型保存位置
The output will be saved to /home/os/AI/CenterNet-master/src/lib/../../exp/ctdet/helmet
heads {'hm': 2, 'wh': 2, 'reg': 2}
Namespace(K=100, aggr_weight=0.0, agnostic_ex=False, arch='res_18', aug_ddd=0.5, aug_rot=0, batch_size=32, cat_spec_wh=False, center_thresh=0.1, chunk_sizes=[16, 16], data_dir='/home/os/AI/CenterNet-master/src/lib/../../data', dataset='helmet', debug=0, debug_dir='/home/os/AI/CenterNet-master/src/lib/../../exp/ctdet/helmet/debug', debugger_theme='white', demo='', dense_hp=False, dense_wh=False, dep_weight=1, dim_weight=1, down_ratio=4, eval_oracle_dep=False, eval_oracle_hm=False, eval_oracle_hmhp=False, eval_oracle_hp_offset=False, eval_oracle_kps=False, eval_oracle_offset=False, eval_oracle_wh=False, exp_dir='/home/os/AI/CenterNet-master/src/lib/../../exp/ctdet', exp_id='helmet', fix_res=True, flip=0.5, flip_test=False, gpus=[0, 1], gpus_str='0,1', head_conv=64, heads={'hm': 2, 'wh': 2, 'reg': 2}, hide_data_time=False, hm_hp=True, hm_hp_weight=1, hm_weight=1, hp_weight=1, input_h=608, input_res=608, input_w=608, keep_res=False, kitti_split='3dop', load_model='', lr=0.000125, lr_step=[90, 120], master_batch_size=16, mean=array([[[0.5684546 , 0.5362932 , 0.49678355]]], dtype=float32), metric='loss', mse_loss=False, nms=False, no_color_aug=False, norm_wh=False, not_cuda_benchmark=False, not_hm_hp=False, not_prefetch_test=False, not_rand_crop=False, not_reg_bbox=False, not_reg_hp_offset=False, not_reg_offset=False, num_classes=2, num_epochs=200, num_iters=-1, num_stacks=1, num_workers=4, off_weight=1, output_h=152, output_res=152, output_w=152, pad=31, peak_thresh=0.2, print_iter=0, rect_mask=False, reg_bbox=True, reg_hp_offset=True, reg_loss='l1', reg_offset=True, resume=False, root_dir='/home/os/AI/CenterNet-master/src/lib/../..', rot_weight=1, rotate=0, save_all=False, save_dir='/home/os/AI/CenterNet-master/src/lib/../../exp/ctdet/helmet', scale=0.4, scores_thresh=0.1,seed=317, shift=0.1, std=array([[[0.25609234, 0.25254378, 0.26302198]]], dtype=float32), task='ctdet', test=False, test_scales=[1.0], trainval=False,val_intervals=5, vis_thresh=0.3, wh_weight=0.1)
Creating model...
#首次运行会下载预训练网络
Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /home/os/.cache/torch/checkpoints/resnet18-5c106cde.pth
ctdet/helmet |# | train: [1][0/29]|Tot: 0:00:07 |ETA: 0:00:00 |loss 16.9254 |hm_loss 14.2204 |wh_loss 22.2494 |off_loss 0.4800 |Da
ctdet/helmet |## | train: [1][1/29]|Tot: 0:00:07 |ETA: 0:03:38 |loss 17.0506 |hm_loss 14.1034 |wh_loss 24.5142 |off_loss 0.4958 |Da
ctdet/helmet |### | train: [1][2/29]|Tot: 0:00:08 |ETA: 0:01:50 |loss 13.7841 |hm_loss 11.0714 |wh_loss 22.0557 |off_loss 0.5072 |Da
ctdet/helmet |#### | train: [1][3/29]|Tot: 0:00:08 |ETA: 0:01:13 |loss 12.3588 |hm_loss 9.5388 |wh_loss 23.1410 |off_loss 0.5058 |Dat
ctdet/helmet |##### | train: [1][4/29]|Tot: 0:00:08 |ETA: 0:00:55 |loss 11.4782 |hm_loss 8.5058 |wh_loss 24.6930 |off_loss 0.5031 |Dat
ctdet/helmet |###### | train: [1][5/29]|Tot: 0:00:09 |ETA: 0:00:44 |loss 10.6189 |hm_loss 7.7164 |wh_loss 24.0185 |off_loss 0.5007 |Dat
ctdet/helmet |####### | train: [1][6/29]|Tot: 0:00:09 |ETA: 0:00:36 |loss 9.9733 |hm_loss 7.1346 |wh_loss 23.4322 |off_loss 0.4955 |Data
ctdet/helmet |######## | train: [1][7/29]|To
Creating model...
=> loading pretrained model https://download.pytorch.org/models/resnet18-5c106cde.pth
Setting up data...
----- val
#记载json数据
==> initializing coco 2017 val data.
annot_path: /home/os/AI/CenterNet-master/src/lib/../../data/helmet/annotations/instances_val2019.json
loading annotations into memory...
Done (t=0.00s)
creating index...
index created!
Loaded val 238 samples
----- train
==> initializing coco 2017 train data.
annot_path: /home/os/AI/CenterNet-master/src/lib/../../data/helmet/annotations/instances_train2019.json
loading annotations into memory...
Done (t=0.01s)
creating index...
index created!
Loaded train 950 samples
Starting training...
epoch: 1
train_loader:
ctdet/helmet/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.
warnings.warn(warning.format(ret))
/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
warnings.warn('Was asked to gather along dimension 0, but all '
开始正式训练
ctdet/helmet |################################| train: [1][28/29]|Tot: 0:00:20 |ETA: 0:00:01 |loss 6.9206 |hm_loss 4.0677 |wh_loss 24.2042 |off_loss 0.4325 |Data 0.167s(0.326s) |Net 0.691s
生成后的模型保存在/exp/ctdet/helmet目录中
(certernet) os@os-Super-Server:~/AI/CenterNet-master/exp/ctdet/helmet$ ls -all
total 775772
drwxrwxr-x 27 os os 4096 7月 13 21:03 .
drwxrwxr-x 4 os os 4096 7月 13 16:02 ..
-rw-rw-r-- 1 os os 237812350 7月 13 17:58 model_120.pth
-rw-rw-r-- 1 os os 237812350 7月 13 17:52 model_90.pth
-rw-rw-r-- 1 os os 80831093 7月 13 17:45 model_best.pth
-rw-rw-r-- 1 os os 237811989 7月 13 20:31 model_last.pth
-rw-rw-r-- 1 os os 2256 7月 13 21:03 opt.txt
5.测试数据
验证测试样本，结果如下；AP和AR很低，因为我的数据只是测试数据，数据量只有一百张左右
#测试单张图片
python demo.py ctdet --demo ../data//helmet/test2019/07022055230001959.jpg --arch res_18 --load_model ../exp/ctdet/helmet/model_best.pth
#测试所有测试集
python test.py --exp_id helmet --not_prefetch_test ctdet --arch res_18 --load_model ../exp/ctdet/helmet/model_best.pth
...
DONE (t=0.25s).
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.461
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.787
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.446
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.161
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.463
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.625
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.224
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.502
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.508
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.226
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.518
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.663
6.报错处理
- int(self.cat_ids[ann[‘category_id’]]) KeyError: 0
file_name: 06291430330001744.jpg /home/os/AI/CenterNet-master/src/lib/../../data/helmet/images/06291430330001744.jpg train
Traceback (most recent call last)
File "main.py", line 104, in
main(opt)
File "main.py", line 72, in main
log_dict_train, _ = trainer.train(epoch, train_loader)
File "/home/os/AI/CenterNet-master/src/lib/trains/base_trainer.py", line 119, in train
return self.run_epoch('train', epoch, data_loader)
File "/home/os/AI/CenterNet-master/src/lib/trains/base_trainer.py", line 61, in run_epoch
for iter_id, batch in enumerate(data_loader):
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 819, in __next__
return self._process_data(data)
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 846, in _process_data
data.reraise()
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/_utils.py", line 369, in reraise
raise self.exc_type(msg)
KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/utils/data/_utils/worker.py", line 178, in _worker_loop
data = fetcher.fetch(index)
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/usr/local/anaconda3/envs/certernet/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 44, in
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/home/os/AI/CenterNet-master/src/lib/datasets/sample/ctdet.py", line 111, in __getitem__
cls_id = int(self.cat_ids[ann['category_id']])
KeyError: 0
解决方法
self.class_name = [
'__background__', 'nothat', 'wearhat']
self._valid_ids = [1, 2 ]
修改为：
self.class_name = [
'__background__', 'nothat', 'wearhat']
self._valid_ids = [ 0, 1, 2 ]