交通标志识别
本项目使用YOLO
模型，并在对数字信号灯进行数字识别时采用opencv算法。
环境安装
所需环境 python =3.7.11 torch==1.2.00
使用
pip install -r requirements.txt
安装所需的包。
文件下载
训练所需的预训练权重可在百度网盘中下载。
链接：https://pan.baidu.com/s/1gKmRdwpQ05fMu1H-mi38zg 提取码：1234
作者训练结果可在下方链接中下载。
链接：https://pan.baidu.com/s/1cLSoWbra612Ezx1EsqOFGQ 提取码： 1234
训练过程
1.数据集的准备
本文使用VOC格式进行训练，训练前需要自己制作好数据集，
训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。
训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
2.数据集的处理
在完成数据集的摆放之后，我们需要利用voc_annotation.py获得训练用的2007_train.txt和2007_val.txt。
修改voc_annotation.py里面的参数。第一次训练可以仅修改classes_path，classes_path用于指向检测类别所对应的txt。
训练自己的数据集时，可以自己建立一个cls_classes.txt，里面写自己所需要区分的类别。
model_data/cls_classes.txt文件内容为：
左转红灯
左转绿灯
...
其中内容也可以换成自己需要的。
3. 开始网络训练
训练的参数较多，均在train.py中，大家可以在下载库后仔细看注释，其中最重要的部分依然是train.py里的classes_path。
classes_path用于指向检测类别所对应的txt，这个txt和voc_annotation.py里面的txt一样！训练自己的数据集必须要修改！
修改完classes_path后就可以运行train.py开始训练了，在训练多个epoch后，权值会生成在logs文件夹中。
4. 训练结果预测
训练结果预测需要用到两个文件，分别是yolo.py和predict.py。在yolo.py里面修改model_path以及classes_path。
model_path指向训练好的权值文件，在logs文件夹里。
classes_path指向检测类别所对应的txt。
完成修改后就可以运行predict.py进行检测了。运行后输入图片路径即可检测。
后处理
- 由于本项目不仅要对红绿灯进行识别，还要对倒计时识别，先采用CNN网络预先对数码管数据集进行训练。然后采用OpenCV对第一步预测出来的结果进行切割，然后把切割出来的图像进行二值化，再进行识别。
预测过程
在yolo.py文件里面，在如下部分修改model_path和classes_path使其对应训练好的文件；model_path对应logs文件夹下面的权值文件，classes_path是model_path对应分的类。
_defaults = {
#--------------------------------------------------------------------------#
# 使用自己训练好的模型进行预测一定要修改model_path和classes_path！
# model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
# 如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
#--------------------------------------------------------------------------#
"model_path" : 'model_data/yolo_weights.pth',
"classes_path" : 'model_data/coco_classes.txt',
#---------------------------------------------------------------------#
# anchors_path代表先验框对应的txt文件，一般不修改。
# anchors_mask用于帮助代码找到对应的先验框，一般不修改。
#---------------------------------------------------------------------#
"anchors_path" : 'model_data/yolo_anchors.txt',
"anchors_mask" : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],
#---------------------------------------------------------------------#
# 输入图片的大小，必须为32的倍数。
#---------------------------------------------------------------------#
"input_shape" : [416, 416],
#---------------------------------------------------------------------#
# 只有得分大于置信度的预测框会被保留下来
#---------------------------------------------------------------------#
"confidence" : 0.5,
#---------------------------------------------------------------------#
# 非极大抑制所用到的nms_iou大小
#---------------------------------------------------------------------#
"nms_iou" : 0.3,
#---------------------------------------------------------------------#
# 该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，
# 在多次测试后，发现关闭letterbox_image直接resize的效果更好
#---------------------------------------------------------------------#
"letterbox_image" : False,
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
"cuda" : True,
}
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
setattr(self, name, value)
#---------------------------------------------------#
# 获得种类和先验框的数量
#---------------------------------------------------#
self.class_names, self.num_classes = get_classes(self.classes_path)
self.anchors, self.num_anchors = get_anchors(self.anchors_path)
self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
self.generate()
#---------------------------------------------------#
# 生成模型
#---------------------------------------------------#
def generate(self):
#---------------------------------------------------#
# 建立yolo模型，载入yolo模型的权重
#---------------------------------------------------#
self.net = YoloBody(self.anchors_mask, self.num_classes)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = self.net.eval()
print('{} model, anchors, and classes loaded.'.format(self.model_path))
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image, imgname):
#---------------------------------------------------#
# 计算输入图片的高和宽
#---------------------------------------------------#
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
# width = 16
# w = image.size[0]
# h = image.size[1]
# w += 2 * width
# h += 2 * width
# img_new = Image.new('RGB', (w, h), (255, 255, 255))
# img_new.paste(image, (width, width))
# images = img_new
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
if results[0] is None:
return image
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
#---------------------------------------------------------#
# 设置字体与边框厚度
#---------------------------------------------------------#
#font = ImageFont.truetype(font='model_data/simhei.ttf',size=20)
font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[0] + 0.5).astype('int32'))
thickness = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
#---------------------------------------------------------#
# 图像绘制
#---------------------------------------------------------#
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = top_conf[i]
top, left, bottom, right = box
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(image.size[1], np.floor(bottom).astype('int32'))
right = min(image.size[0], np.floor(right).astype('int32'))
label = '{}'.format(predicted_class)
# 判断倒计时灯颜色， 进行预处理
if predicted_class == '红色倒计时':
np_image = np.array(image)
np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
roi = np_image[top:bottom, left:right]
# roi = np_image[top-1:bottom+1, left-1:right+1]
cv2.imwrite('./analysis/%s_roi.jpg'%imgname, roi)
# gray_image = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
gray_image = tomygray(roi, predicted_class)
cv2.imwrite('./analysis/%s_gray_img.jpg'%imgname, gray_image)
predict_text = main_test(gray_image, imgname)
# label = '{} {}'.format(predicted_class, predict_text)
label = '倒计时：{}'.format(predict_text)
elif predicted_class == '绿色倒计时':
np_image = np.array(image)
roi = np_image[top:bottom, left:right]
# roi = np_image[top-1:bottom+1, left-1:right+1]
gray_image = tomygray(roi, predicted_class)
predict_text = main_test(gray_image, imgname)
# label = '{} {}'.format(predicted_class, predict_text)
label = '倒计时：{}'.format(predict_text)
elif predicted_class == '黄色倒计时':
np_image = np.array(image)
roi = np_image[top:bottom, left:right]
# roi = np_image[top-1:bottom+1, left-1:right+1]
gray_image = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
predict_text = main_test(gray_image, imgname)
# label = '{} {}'.format(predicted_class, predict_text)
label = '倒计时：{}'.format(predict_text)
预测结果
