TensorFlow2.x-YOLOv3训练自定义数据库
本文记录如何用TensorFlow2.x-YOLOv3训练自定义数据库。训练使用TensorFlow-cpu 2.2，在Intel(R) Core(TM) i7-8550U CPU 8G RAM的笔记本中完成，训练artelab图集427张图的时间约为14个小时，供参考。
参考链接
tensorflow-yolov3 基于tensorflow-gpu==1.11.0版本实现的yolo3版本
TensorFlow2.0-Examples基于tensorflow-gpu==2.0版本实现的yolo3版本
YunYang1994/tensorflow-yolov3 训练自己的数据集
artelab本文训练图集下载链接，图集内容为条形码。本文以训练该图集为例，来说明如何使用TensorFlow2.x-YOLOv3训练自定义的数据库。
向原作者致敬，非常感谢！
运行环境
Win10 + Python 3.6
TensorFlow-cpu 2.2
Keras 2.3.1
训练前准备工作
1. Github中下载TensorFlow2.0-Examples
下载后，请尝试运行TensorFlow2.0-Examples/4-Object_Detection/YOLOV3/test.py，若可以正常运行，则表示当前电脑用TensorFlow2.x-YOLOv3训练自定义的数据库的基础条件已具备。若有需要，请参考TensorFlow2.x-YOLOv3在Win10中运行Demo。
2. Github中下载labelImg并确认其可以正常运行。如有需要，请参考Install labelImg on Win10 + python3.6
3. 下载artelab训练图集，下图为artLab图集中的一张，或准备其他训练图集。
TensorFlow2.x-YOLOv3训练自定义数据库
1. TensorFlow2.x-YOLOv3文件结构
Github下载得到的TensorFlow2.x-YOLOv3文件结构为：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
D:\GITHUB\TENSORFLOW2.0-EXAMPLES\4-OBJECT_DETECTION\YOLOV3 │ image_demo.py │ README.md │ test.py │ train.py │ video_demo.py │ yolov3.weights │ ├─core │ │ backbone.py │ │ common.py │ │ config.py │ │ dataset.py │ │ utils.py │ │ yolov3.py │ │ __init__.py │ │ │ └─__pycache__ │ ├─data │ ├─anchors │ │ basline_anchors.txt │ │ │ ├─classes │ │ coco.names │ │ yymnist.names │ │ │ └─dataset │ yymnist_test.txt │ └─docs kite.jpg requirements.txt road.mp4 timg.jpg |
在以上结构中，我们需要更改YOLOV3/data文件夹下的内容，以及YOLOV3/core/config.py中的部分代码。接下来，我们逐一说明。
2. YOLOV3/data准备
将上述TensorFlow2.x-YOLOv3单独拷贝出来，改造成我们所需要的文件结构。改造完成的TensorFlow2.x-YOLOV3/data的文件结构图如下：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
D:\PYTHON\TENSORFLOW2.X_YOLOV3\DATA ├─anchors │ basline_anchors.txt │ ├─classes │ barcode.names │ └─dataset │ ├─convert.py │ ├─imageset │ ├─test.txt │ ├─train.txt │ └─yolotxt |
将原有的文件结构改造为上述结构，我们需要新增或修改以下内容：
data/classes：删除原有的coco.names和yymnist.names，新增barcode.names。本文示例中仅需要检测barcode一类目标，因此 barcode.names内容为：
|
1 |
barcode |
data/dataset: 删除原有的yymnist_test.txt，新增以下内容：
data/dataset/imageset: 待训练的图片集，所有用于训练的图片均放在此路径下。
data/dataset/yolotxt: 用labelImg标注的训练图集中的目标信息txt文件所在路径，其中的文件数量与待训练图片集数量相同，文件名称与待训练图片集的图片一一对应。
data/dataset/convert.py: 用于将labelImg标注的目标信息YOLO格式的数据转换为TensorFlow2.x-YOLOv3格式的数据。转换的原理以及原因请参阅本文小节【YOLO数据转换为TensorFlow2.x-YOLOv3数据】。
data/dataset: 该文件夹中test.txt和train.txt由convert.py根据data/dataset/yolotxt生成。默认情况下，将待训练图片集随机分类，取75%的图片作为训练集，目标信息存放在train.txt中；另外25%作为测试集，目标信息存放在test.txt中。若需要更改比例，需要在convert.py中修改相应参数。
3. 修改YOLOV3/core/config.py
打开config.py，将其中的
|
1 2 3 4 5 |
__C.YOLO.CLASSES = "./data/classes/coco.names" __C.TRAIN.ANNOT_PATH = "./data/dataset/yymnist_train.txt" __C.TEST.ANNOT_PATH = "./data/dataset/yymnist_test.txt" |
分别修改为：
|
1 2 3 4 5 |
__C.YOLO.CLASSES = "./data/classes/barcode.names" __C.TRAIN.ANNOT_PATH = "./data/dataset/train.txt" __C.TEST.ANNOT_PATH = "./data/dataset/test.txt" |
上述三行代码并不在一起，这里为方便描述，写在一起。分别修改的是目标类型文件名，训练和测试数据信息路径与文件名。上述文件均使用的是相对路径，相对于TensorFlow2.x-YOLOv3原文件结构中的TensorFlow2.x-YOLOv3/train.py的路径。
4. 开始训练
以上内容修改完成后，在TensorFlow2.x-YOLOv3/train.py路径下运行train.py文件。正常开始训练的界面如下
|
1 2 3 4 5 6 7 |
D:\Python\Tensorflow2.x_YOLOV3>train.py 2020-07-12 08:40:40.917962: I tensorflow/core/platform/cpu_feature_guard.cc:143] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 2020-07-12 08:40:40.950656: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x232384ac3d0 initialized for platform Host (this does not guarantee that XLA will be used). Devices: 2020-07-12 08:40:40.957342: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version => STEP 1 lr: 0.001000 giou_loss: 3.99 conf_loss: 1952.01 prob_loss: 1.98 total_loss: 1957.97 => STEP 2 lr: 0.000011 giou_loss: 4.28 conf_loss: 2022.44 prob_loss: 2.30 total_loss: 2029.02 => STEP 3 lr: 0.000017 giou_loss: 4.26 conf_loss: 2002.26 prob_loss: 2.25 total_loss: 2008.77 |
我的电脑配置是Intel(R) Core(TM) i7-8550U CPU 8G RAM，训练artelab图集427张图的时间约为14个小时，供大家参考。
5.训练完成后，在train.py所在路径下运行tensorboard –logdir ./data/log，可以查看训练过程
|
1 2 3 |
D:\Python\Tensorflow2.x_YOLOV3>tensorboard --logdir ./data/log Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all TensorBoard 2.2.2 at http://localhost:6006/ (Press CTRL+C to quit) |
此时，在浏览器中输入http://localhost:6006/，可以查看图形化的训练过程。
TensorFlow2.x-YOLOv3自定义数据库训练模型测试
在Tensorflow2.x_YOLOV3路径下，运行python test.py，即可使用当前训练的模型来检测我们图集中的测试集。
|
1 2 3 4 5 6 7 8 9 10 11 12 |
D:\Python\Tensorflow2.x_YOLOV3>python test.py 2020-06-01 08:01:16.603918: I tensorflow/core/platform/cpu_feature_guard.cc:143] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 2020-06-01 08:01:16.691544: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x1734bd4b540 initialized for platform Host (this does not guarantee that XLA will be used). Devices: 2020-06-01 08:01:16.739968: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version => ground truth of 05102009137.jpg: barcode 532 470 1450 1056 => predict result of 05102009137.jpg: barcode 0.7751 542 382 1474 1119 => ground truth of 05102009145.jpg: barcode 320 149 1248 1025 => predict result of 05102009145.jpg: barcode 0.8194 230 127 1319 1070 |
YOLO数据转换为TensorFlow2.x-YOLOv3数据
在Install labelImg on Win10 + python3.6中提到过，labelImg标注的YOLO数据意义为：
图片标注完成后，每张图片均会生成一张与图像名称相同的txt文件，其内容及其意义如下：
|
1 2 3 |
0 0.569583 0.474375 0.562500 0.295000 class_id x/imgWidth y/imgHeight w/imgWidth h/imgHeight |
在上述数据中，相应变量的意义如下：
x/y: 标注区域中心点坐标
w/h: 标注区域宽与高
imgWidth/imgHeight : 源图像的宽度与高度。
而TensorFlow2.x-YOLOv3需要将所有图片目标信息集中在一个txt文档中，txt文档中的一行代表一张图片的信息，一张图片中可以有多个目标信息，每行数据的意义如下：
|
1 2 3 4 |
xxx/xxx.jpg 18.19,6.32,424.13,421.83,20 323.86,2.65,640.0,421.94,20 xxx/xxx.jpg 48,240,195,371,11 8,12,352,498,14 # image_path x_min, y_min, x_max, y_max, class_id x_min, y_min ,..., class_id # make sure that x_max < width and y_max < height |
正因为有了上述供需的差异，所以我们需要依靠convert.py来将labelImg标注的YOLO数据转换为TensorFlow2.x-YOLOv3所需要的数据。该文件的原理可以参考这里。
我用的文件基本与原博客相同，做了一些简单的修改，请参阅下面的代码。
关于convert.py，有以下几个注意点：
1. 变量source_img_path_related_to_train_py与source_img_path
这两个变量指向的其实是同一个文件夹，即待训练图片所在文件夹imageset。但它们都是相对路径，source_img_path_related_to_train_py是imageset相对于train.py所在位置的路径，而source_img_path则是相对于convert.py所在位置的路径。
2. 变量MAX_ID
该变量值应该与待训练图片的类别数量相同。
3. 变量train_scale
该变量设置图像集中用作训练集的百分比，默认为0.75，若需要修改，修改此变量即可。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import os import re import shutil import cv2 import random def extract_content(content_): # 注意，一开始用的第一种，结果只有一行的情况没有被提取出来，要去掉后面的\n，谨记 # content_extract = re.findall('(.*?) (.*?) (.*?) (.*?) (.*?)\n', content) content_extract_ = re.findall('(\d+.?\d*) (\d+.?\d*) (\d+.?\d*) (\d+.?\d*) (\d+.?\d*)', content_) return content_extract_ if __name__ == '__main__': # 记得路径尾部加“/”，不然调用join方法是它会用“\”替代，那样不好，容易造成转义字符问题。 # ../表示上一层路径 # 该路径会写入train.txt和test.txt，tensorflow-yolov3训练时使用，相对于train.py的路径 source_img_path_related_to_train_py = './data/dataset/imageset/' #转换格式时所使用的路径，相对于convert.py的路径 source_img_path = './imageset/' #用labelImg标注的yolo格式的目标信息文件路径。与训练图集中的图像一一对应，名称相同，类型为txt文档。 source_txt_path = './yolotxt/' #yolo格式转换为tensorflow2-yolo3格式后的文件存储路径 target_txt_path = './' # train.txt:记录训练用图像路径以及相应图像中目标相关信息 # test.txt:记录测试用图像路径以及相应图像中目标相关信息 # 打开上述两个文件准备写入 train_file = open(target_txt_path + 'train.txt', 'w', encoding='utf-8') test_file = open(target_txt_path + 'test.txt', 'w', encoding='utf-8') # 创建写入内容字符串变量 train_file_content = '' test_file_content = '' # 该参数与实际训练类别数量相同 MAX_ID = 1 # 读取source_txt_path路径下所有文件（包括子文件夹下文件） filenames = os.listdir(source_txt_path) # 打开labelImg标注的yolo格式的文件，提取其中数字并将内容重构后写入新文件 for filename in filenames: # 打开文件： with open(os.path.join(source_txt_path, filename), 'r', encoding='utf-8') as f: # 读取文件内容 content = f.read() # 提取数据 content_extract = extract_content(content) # 获取当前图片分辨率信息（这样不论图片尺寸多少都能成功转换）（re.findall()返回的是列表，需要将它转换成字符串） # 读取图片 tmpfilename = '{}{}.jpg'.format(source_img_path, ''.join(re.findall('(.*?).txt', filename))) print(tmpfilename) img = cv2.imread(tmpfilename) # 获取图片分辨率 img_width = img.shape[1] img_height = img.shape[0] # 创建单行写入字符串的路径头字符串：图片路径+图片名称 path_str = source_img_path_related_to_train_py + os.path.splitext(filename)[0] + '.jpg' print(path_str) # 创建单行写入字符串的目标坐标字符串 obj_strs = '' # 将数据格式从相对坐标转换成绝对坐标 for obj_str in content_extract: # 将元组字符串转换成列表数字 object_evar = list(map(eval, obj_str)) # 映射变量 class_id = object_evar[0] x, y = object_evar[1] * img_width, object_evar[2] * img_height w, h = object_evar[3] * img_width, object_evar[4] * img_height # 判断数据是否超出限制（数据清洗）（包括清洗超限坐标和错误class_id） if class_id >= MAX_ID \ or round(x - w / 2) < 0 \ or round(x + w / 2) > img_width \ or round(x - w / 2) >= round(x + w / 2) \ or round(y - h / 2) < 0 \ or round(y + h / 2) > img_height \ or round(y - h / 2) >= round(y + h / 2): print('错误标注：') print(filename) print(object_evar) print('[{}, {}, {}, {}, {}]'.format(round(x - w / 2), round(y - h / 2), round(x + w / 2), round(y + h / 2), class_id)) continue # 将映射变量格式化后加入到obj_strs中： obj_strs += ' {},{},{},{},{}'.format(round(x - w / 2), round(y - h / 2), round(x + w / 2), round(y + h / 2), class_id) # 拆分训练集和测试集 # 训练集占比 train_scale = 0.75 # 设置随机概率 proba = random.random() # 如果该张图片经过数据清洗后没有目标，则跳过，不将其加入到train.txt和test.txt文件中 if obj_strs == '': print('空文件：{}'.format(filename)) print('content：{}'.format(content)) cv2.imwrite('null_img\\{}.jpg'.format(''.join(re.findall('(.*?).txt', filename))), img) print('将图片拷贝到“空文件”文件夹') continue else: write_strs = path_str + obj_strs print(write_strs) # 判断该写入哪个文件 if proba < train_scale: train_file_content += write_strs + '\n' else: test_file_content += write_strs + '\n' # 将两个即将写入的内容去除首位的无效字符（如空格，换行符，制表符，回车符） train_file_content = train_file_content.strip() test_file_content = test_file_content.strip() train_file.write(train_file_content) test_file.write(test_file_content) train_file.close() test_file.close() |
如何在自己的程序中调用训练模型
按照上述步骤训练完成后，在train.py路径下会发现“yolov3.index”和”yolov3.data-00000-of-00001″这两个文件，它们就是我们训练得到的模型。借助Tensorflow2.x_YOLOV3中原有的core和data中的部分文件，参考test.py，构造文件夹可以调用我们训练得到的模型。如下文件结构所示，your_code.py为准备调用训练模型的文件。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
D:\PYTHON\YOLOV3 │ │ your_code.py │ ├─core │ │ backbone.py │ │ common.py │ │ config.py │ │ dataset.py │ │ utils.py │ │ yolov3.py │ │ __init__.py │ │ │ └─__pycache__ │ backbone.cpython-36.pyc │ common.cpython-36.pyc │ config.cpython-36.pyc │ dataset.cpython-36.pyc │ utils.cpython-36.pyc │ yolov3.cpython-36.pyc │ __init__.cpython-36.pyc │ ├─data │ ├─anchors │ │ basline_anchors.txt │ │ │ └─classes │ barcode.names │ └─model yolov3.data-00000-of-00001 yolov3.index |
在your_code.py文件中准备以下代码，即可调用训练得到的模型进行目标检测。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import cv2 import numpy as np import tensorflow as tf from tensorflow.python.keras import layers import core.utils as utils from core.config import cfg from core.yolov3 import YOLOv3, decode INPUT_SIZE = 416 def ini_model(): NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES)) CLASSES = utils.read_class_names(cfg.YOLO.CLASSES) # Build Model input_layer = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3]) feature_maps = YOLOv3(input_layer) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = decode(fm, i) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.load_weights('./model/yolov3') return model model = ini_model() image = cv2.imread("your_test_image_path") #Read the test image image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image_size = image.shape[:2] image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE]) image_data = image_data[np.newaxis, ...].astype(np.float32) pred_bbox = model.predict(image_data) pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox] pred_bbox = tf.concat(pred_bbox, axis=0) bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD) bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms') image = utils.draw_bbox(image, bboxes) cv2.imshow("test image", image) cv2.waitKey(0) |
本文到此结束，感谢阅读，欢迎关注。
