TensorFlow下使用YOLOv1訓練自己的數據集+測試自己的模型

一.?前期準備

環境：（用cpu跑的）

win10 + python3.6.8?+ tensorflow2.4.1+pycharm

ps：本來打算用tensflow-gpu? 1.4.2運行的，但是該代碼是2.xx版本的tensflow,需要安裝tensflow-gpu? 2.x.版本，以及cuda 11.0版本的 ，由于本人電腦比較差，最高cuda安裝到8.0版本，又不想升級顯卡驅動，所以用cpu跑的。

YOLOv1代碼：https://e.coding.net/xucancan1/yolov1_tensorflow/xucc1617608769.git

1.1準備自己要訓練的圖片

?first?在data文件夾中創建一個myDataSet文件夾來保存爬取的圖片，修改 pachong.py文件的保存地址，運行pachong.py文件，爬蟲完后然后從myDataSet文件夾中挑選出符合自己數據集（比如我想訓練只檢測手機，所以挑選帶手機的圖片）的圖片，刪除不符合的圖片

second 再接著 運行SelectSize.py文件，繼續篩選文件

last? 最后運行graph_rename.py文件，給上述篩選的文件命名，得到最終的原圖片

我本人在網上爬取的圖片：爬取代碼以及其他代碼如下：

GitHub - xucancan1617608769/yolov1

1.2 數據集準備

在Yolo_tensorflow\data\pascal_voc\VOCdevkit\VOC2007路徑下，創建如下文件夾（ps:我的文件夾中已經創好了）

?將1.1中data文件夾中的圖片全部copy到此路徑下的JPEGImages文件夾中,然后在ImageSets文夾下創建Main文件夾，接著用labelImg工具對JPEGImages文件夾的圖片進行數據標記保存到Annotations文件夾下

?工具：LabelImg? ?鏈接：鏈接：https://pan.baidu.com/s/1qwPnIx-T_-Kl1CawtE4M9g?
提取碼：8cb2

labelImg工具的使用參考鏈接：使用LabelImg標注圖片 - AiFly - 博客園

?像這樣：? ? ??

運行conert_to_txt.py文件生成ImageSets/Main/文件夾下的4個文件。（layout文件夾不用管）

??? 像這樣：

運行完后Main文件夾會生成yolov1所需的train.txt,val.txt,test.txt,trainval.txt

至此，自己要訓練的數據集已準備完畢

二.修改配置

①修改?..\yolo_tensorflow\yolo\config.py文件? ?如下

CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus','car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse','motorbike', 'person', 'pottedplant', 'sheep', 'sofa','train', 'tvmonitor']

====>>CLASSES = ['phone', 'cat', 'dog', 'sheep'] # 修改為自己的類別

?② ..\yolo_tensorflow\utils\pascal_voc.py文件

#第30行labels = np.zeros((self.batch_size, self.cell_size, self.cell_size, 9)) # 25修改為5+類別數#第129行label = np.zeros((self.cell_size, self.cell_size, 9)) # 修改為5+類別數

開始訓練，運行?train.py文件。

三 .保存權重文件? 測試?

① 將..\yolo_tensorflow\data\pascal_voc\output\2021_09_12_10_47\目錄下（2021_09_12_10_47文件夾在模型開始訓練時自動生成）保存的最后的模型文件（下圖1）復制到..\yolo_tensorflow\data\weights\目錄下（下圖2），并在名稱中間添加“.ckpt”重新命名。

? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?上圖1

? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?上圖2

② 修改..\yolo_tensorflow\test.py文件。

將待測試圖片放入..\yolo_tensorflow-master\test\目錄下。

# 在189行的默認權重文件改為自己訓練好權重的模型 def main():parser = argparse.ArgumentParser()# parser.add_argument('--weights', default="YOLO_small.ckpt", type=str)parser.add_argument('--weights', default="yolo-15000.ckpt", type=str)# 將207行需要測試的圖片名稱修改為自己的圖片名稱# detect from image fileimname = 'test/000004.jpg'detector.image_detector(imname)

③?運行test.py文件進行測試

結果：

?四.結語

整個流程按照步驟走了一遍，中間需要安裝一些第三方安裝包，比如tensorflow2.x.x版本，以及cv2,tf_sim,最后結果不是很理想，定位存在很大誤差，以后再更新改進。

總結

以上是生活随笔為你收集整理的Yolov1-手把手用自己的数据集训练自己的模型的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。