交通標識の認識
このプロジェクトでは、YOLO
モデルを使用し、デジタル信号灯のデジタル認識に opencv アルゴリズムを使用します。
環境のインストール
必要環境 python=3.7.11 torch==1.2.00
使用
pip install -r requirements.txt
必要なパッケージをインストールします。
ドキュメントをダウンロード
トレーニングに必要な事前トレーニング済みの重みは、Baidu Cloud Disk からダウンロードできます。
リンク: 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 のみを変更できます。これは、検出カテゴリに対応するテキストを指すために使用されます。
独自のデータセットをトレーニングする場合は、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 を実行してトレーニングを開始できます。複数のエポックのトレーニング後、重みがログ フォルダーに生成されます。
4. トレーニング結果の予測
トレーニング結果の予測には、yolo.py と detect.py という 2 つのファイルを使用する必要があります。yolo.pyのmodel_pathとclasses_pathを変更します。
model_path は、ログ フォルダー内のトレーニングされた重みファイルを指します。
classes_path は、検出カテゴリに対応する txt を指します。
変更が完了したら、predict.py を実行して検出できます。実行後、検出するイメージのパスを入力します。
後処理
- このプロジェクトでは信号機を認識するだけでなく、カウントダウンも認識する必要があるため、最初に CNN ネットワークを使用してデジタル チューブ データ セットを事前トレーニングします。次に、最初のステップで予測された結果を OpenCV を使用して切り取り、切り取った画像を 2 値化して認識します。
予測プロセス
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)
予測結果
