附录2-tensorflow目标检测

news2024/9/21 0:34:13

源码来自作者Bubbliiiing，我对参考链接的代码略有修改，网盘地址

链接：百度网盘请输入提取码提取码：dvb1

1 参考链接

2 环境

3 数据集准备

3.1 VOCdevkit/VOC2007

3.2 model_data/voc_classes.txt

3.3 voc_annotation.py

4 训练 train.py

5 训练结果

6 预测

7 其他

7.1 多线程训练

7.2 二次训练

7.3 学习速率

1 参考链接

源码地址 GitHub - bubbliiiing/yolo3-tf2: 这是一个yolo3-tf2的源码，可以用于训练自己的模型。

博客地址睿智的目标检测51——Tensorflow2搭建yolo3目标检测平台_Bubbliiiing的博客-CSDN博客_yolo3

视频地址睿智的目标检测51——Tensorflow2搭建yolo3目标检测平台_Bubbliiiing的博客-CSDN博客_yolo3

2 环境

系统 Linux
显卡 NVIDIA GeForce RTX 3060
CUDA 11.1
CUDNN 无 (cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2与cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2都查不到)

python版本3.6，环境如下

我直接用这个whl装的，tensorflow_gpu-2.6.0-cp36-cp36m-manylinux2010_x86_64.whl

装完之后将keras降到了2.6.0

训练时默认使用GPU资源进行训练

项目放在home下，项目命名为tensorflow_object_detection

3 数据集准备

数据集为877张图像，4分类，其中speedlimit 705个框，crosswalk 174个框，traffclight 154个框，stop 88个框

3.1 VOCdevkit/VOC2007

在项目路径下的VOCdevkit/VOC2007中，将Annotations放入标注的XML文件，JPEGImages放入标注的图片文件(必须是jpg格式的图像，其他格式的不行)

进入ImageSets/Main，删除其中的所有内容

删除项目路径下的 2007_train.txt与2007_val.txt

3.2 model_data/voc_classes.txt

打开项目路径下model_data中的voc_classes.txt

将里面的内容改为自己要训练的类别，顺序无所谓

3.3 voc_annotation.py

不需要改动代码直接运行 voc_annotation.py

运行后会生成这些文件

4 训练 train.py

根据需要修改这里的epoch

然后直接运行就好了，一些warning可以无视掉

在训练开始的时候会给一些提示，可根据这里的提示修改上面的epoch，比如我现在就将epoch设置为569

训练会持续很长事件

5 训练结果

训练结束后会在logs中出现一些文件，我们预测的时候使用 best_epoch_weights.h5 就可以了

我们可以在训练过程中，或者在训练好的loss文件中，查看loss情况

在epoch_loss.txt中可以查看具体的数值

看下面这两个哪个都行

6 预测

我简单改了一下源代码中yolo.py的detect_image方法，目的是拿到预测的信息，而不是直接得到图像

文件名改为了Suyu_yolo.py，下面的predict.py中会进行调用

然后改了一下源码中的predict.py(文件名我改为了Suyu_predict.py)

import time
import cv2
import numpy as np
import tensorflow as tf
from PIL import Image
from Suyu_yolo import YOLO
from utils.utils import get_classes

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

yolo = YOLO()

class_names,num_classes = get_classes('model_data/voc_classes.txt')
img = './img/road344.jpg'
image = Image.open(img)
out_boxes, out_scores, out_classes = yolo.detect_image(image)

result_img = cv2.imread(img)
for i, c in list(enumerate(out_classes)):
    predicted_class = class_names[int(c)]
    box = out_boxes[i]
    score = out_scores[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 = '{} {:.2f}'.format(predicted_class, score)
    print(label)

    cv2.rectangle(result_img,(left,top),(right,bottom),(0,255,0),2)
    cv2.putText(result_img,label,(left,top+5),cv2.FONT_HERSHEY_SIMPLEX,1,(255,0,0),2)

cv2.imshow('result_img',result_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

之后我们将一张图像放在文件夹img中