基于Faster-rcnn的安全帽检测(使用Colab进行训练)
前言
寒假学习了Faster-rcnn算法流程并且阅读了别人复现的代码,感觉收获颇多,但总感觉缺乏实践,因此找了一份开源的数据集,希望能够通过这个项目入门目标检测以及加深对Faster-rcnn的理解。
源码与数据集下载
别人复现的Faster-RCNN源码
源码对应的博客
数据集地址
可以在Colab上进行训练的源码 说明: 我是在这份源码的基础上进行修改的
数据预处理
数据集目录结构
源码是按照VOC
格式的数据进行处理,因此下载安全帽数据集,并把其整理为VOC
格式。
划分训练、验证、测试集
运行步骤
-
首先在
VOC2028
文件夹下新建voc_classes.txt
文件,包含数据集包含的类别的名字。顺序没有关系,但是名字一定要和xml
文件里object
里的name
匹配,person
不能在这里写成people
-
运行
voc_annotation.py
文件划分训练、验证、测试集。
补充说明
在voc_annotation.py
文件中,有多种模式:
0: 干了两件事。
-
- 在
ImageSets
目录产生四个txt文件,包括训练集、验证集、测试集的图像id
,但是不包括后缀。在跑安全帽数据的时候,我记得有些图片后缀是JPG
(大写),导致后面读取数据出错,记得把后缀名都统一为小写的jpg
。
- 在
-
- 根据上一步产生的txt文件生成对应的包含标注信息的txt文件。
格式是:img_path box1 box2...boxn
其中boxi的格式是x1, y1, x2, y2 类别的id
- 根据上一步产生的txt文件生成对应的包含标注信息的txt文件。
第一次运行的mode设为0即可
项目根目录
关于为什么要在前面写这几行代码的:
# 建议在本机上op = 1先生成txt文件再上传到云盘,会快很多
op = 0 # 0 for 本机 1 for colab
data_year = '2028'
if op == 0:
project_root = pathlib.Path(__file__).parent
else:
project_root = pathlib.Path('/content/drive/MyDrive/faster-rcnn-pytorch-master')
在Colab中,但凡涉及到路径的,一定要使用绝对路径,不能用相对路径,不然会出现no file or directory
等错误,所以需要获得项目的根路径,然后后面凡是涉及到xxx_path
的都要和根目录进行拼接。
我发现在Colab上运行这份代码,产生txt文件的过程特别慢,我想大概是程序需要从云盘里一个个文件的去读取所以很慢,所以建议在本机上op = 1先生成txt文件再上传到云盘,会快很多。
数据可视化
我认为碰到任何一个数据集,首先得把标注box标在原图上,观察目标的大小、数目等特点,才方便去调参。
在utils/utils.py
文件,draw_annotation
能够实现把标注信息标在原图上。有两种模式:
batch = True
:批量标注batch = False
:单张标注
def draw_annotation(img_path, annot_path, save_path, classes_path, batch=False):
"""
img_path:原图的目录路径
annot_path:标注文件的路径
save_path:把框画在原图上之后保存的目录路径
"""
font = cv2.FONT_HERSHEY_SIMPLEX
class_names, num_classes = get_classes(classes_path)
if batch:
print('start draw ground truth in batch way')
f = open(annot_path, 'r', encoding='UTF-8')
annot = f.readlines()
num = 0
for line in tqdm(annot):
line = line.split()
img = cv2.imread(line[0])
img_name = os.path.basename(line[0])
img_copy = img.copy()
for box in line[1:]:
box = box.split(',')
cv2.rectangle(img_copy, (int(box[0]), int(box[1])), (int(
box[2]), int(box[3])), color=(0, 255, 0), thickness=3)
label = class_names[int(box[-1])]
cv2.putText(img_copy, label, (int(box[0]), int(box[1])), font,
1.2, (0, 0, 255), 2)
if not os.path.exists(save_path):
os.mkdir(os.path.join(save_path))
cv2.imwrite(os.path.join(save_path, img_name), img_copy)
num += 1
print('draw {} ground truth in batch way done!'.format(num))
else:
img_name = os.path.basename(img_path)
img = cv2.imread(img_path)
# img = img[:,:,::-1]
f = open(annot_path, 'r')
annot = f.readlines()
a = img.copy()
for line in annot:
line = line.split()
annot_name = os.path.basename(line[0])
if annot_name[:-4] == img_name[:-4]:
for box in line[1:]:
box = box.split(',')
cv2.rectangle(a, (int(box[0]), int(box[1])), (int(
box[2]), int(box[3])), color=(0, 0, 255))
cv2.imshow('1', a)
cv2.waitKey(0)
我是放在了和原数据同级目录下,在之后训好模型之后再测试集上把预测的框画在原图上也保存起来,这样真值和预测值都可视化了,就方便观察了。
# 把这个函数跑三次,路径也要相应的改
# 第一次
draw_annotation(img_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\JPEGImages',
annot_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\train.txt',
save_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\Trainval-Ground-truth',
classes_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\voc_classes.txt',
batch=True)
# 第二次
draw_annotation(img_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\JPEGImages',
annot_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\val.txt',
save_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\Trainval-Ground-truth',
classes_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\voc_classes.txt',
batch=True)
# 第三次
draw_annotation(img_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\JPEGImages',
annot_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\test.txt',
save_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\Test-Ground-truth',
classes_path=r'D:\faster-rcnn-pytorch-master\VOCdevkit\VOC2028\voc_classes.txt',
batch=True)
训练
模型
源码实现了vgg16
和resnet50
作为backbone,考虑到resnet50
可能会比较大,因此我只用了vgg16
训练。
首先在nets/vgg16.py
文件里第三行的load_state_dict_from_url
函数修改一下包来源:from torch.hub import load_state_dict_from_url
参数设置
在utils\config.py
文件里进行超参数的设置。
根据第二步数据可视化的结果来看,这个数据集的目标既有大而稀,也有小而密的,因此我在anchor_size
的设置上,相比原论文就会多一些,并且base_size
也要调小一点。
开始训练
建议下载原博主用vgg16为backbone在VOC训练好的模型,放到model_data
,加载这个模型后再训练安全帽数据集,效果会好很多。
新建一个Colaboratory,输入以下命令行即可。
from google.colab import drive
drive.mount('/content/drive')
!pip install yacs
!python drive/MyDrive/faster-rcnn-pytorch-master/train.py
Colab每运行12h就会中断,下次再进行训练时可以使用以下命令:
from google.colab import drive
drive.mount('/content/drive')
!pip install yacs
!python drive/MyDrive/faster-rcnn-pytorch-master/train.py --checkpoint exp/VOC2028/vgg16/checkpoint.pth
获得结果
保存网络参数
在Colab训练到17个Epoch之后就迫不及待的想要测一下结果,测试的话我认为还是在本地会更方便。
可以看到参数文件有1.5G左右大小(原因是不仅保存了网络参数,还保存了优化器、lr_schedule等的参数,目的是中断后重新训练更方便),相信使用Colab的同学都是缺卡的(哭唧唧),这么大的参数我在本地笔记本上用torch.load
显存会爆,所以只能在Colab上读取现有的model_best.pth
文件,然后把网络的参数单独保存为一个pth
文件,这样文件就只有500+M
大小了。
在Colab上用以下命令即可保存网络的参数到项目根目录下的model_best.pth
文件。
from google.colab import drive
drive.mount('/content/drive')
import os
import torch
project_root = r"/content/drive/MyDrive/faster-rcnn-pytorch-master"
model_path = r"exp/VOC2028/vgg16/model_best.pth"
pretrained_dict = torch.load(os.path.join(project_root,model_path), map_location=torch.device('cpu'))['state_dict']
torch.save(pretrained_dict, r"/content/drive/MyDrive/faster-rcnn-pytorch-master/model_best.pth", _use_new_zipfile_serialization=False)
获得map
运行get_map.py
文件可以得到map结果。结果还是不是很好,还需要继续调参。