tiny YOLO v3做缺陷檢測實戰

前言：接觸yolo網絡是在七月份，當時把yolo檢測的論文以及R-CNN系列，SSD等一些論文看了一下，感覺內容很豐富，也嘗試了darknet版本的實現，和yolo v3的實現，在網上也有很多關於上面兩種的實現，這裏就不講了。九月份用tiny-yolo v3做了一個缺陷檢測的實驗，效果出乎意料，準確率和召回率“滿分”！！過了三個月纔想着把以前的實驗總結一下，真不應該。下面從頭開始說明怎麼在自己的數據集上實現tiny-yolo v3，碼字不易，給贊啊，代碼是在別人的基礎上修改的，放到了github上，感謝加星：https://github.com/Eatzhy/tiny-yolov3  有問題歡迎交流，會幫忙解決。
1、對圖像進行轉格式和編隊

因爲實驗使用的缺陷圖像爲DAGM的數據集，大小爲512×512格式爲PNG，而tiny-yolo v3輸入的格式爲jpg，大小爲416×416(也有人說tiny網絡輸入是224×224大小，我們先不管到底哪個，代碼在修改的時候是416)，對圖像轉格式，編隊代碼如下。

將代碼copy下來和缺陷圖像文件夾放到一個路徑，改一下代碼的路徑名就ok。轉換完之後大概是這個樣子（部分圖片）：

# -*- coding: utf-8 -*-
 
'''
將png轉jpg
resize到416
並給圖片編隊
待轉換的圖像存放在data下，程序運行後，data下獲得的是jpg格式
pre_data存放是png格式
'''
import os 
from PIL import Image 
import shutil 
import sys 
  
#創建一個文件，存放原圖
output_dir = 'pre_data' 
if not os.path.exists(output_dir): 
  os.mkdir(output_dir) 
    
def image2png(dataset_dir,type):
    #轉換格式並resize到一定大小
    files = [] 
    image_list = os.listdir(dataset_dir) 
    files = [os.path.join(dataset_dir, _) for _ in image_list] 
    for index,png in enumerate(files): 
        if index > 100000: 
            break
        try:
            sys.stdout.write('\r>>Converting image %d/100000 ' % (index)) 
            sys.stdout.flush() 
            img = Image.open(png)
            img = img.resize((416,416))
            jpg = os.path.splitext(png)[0] + "." + type
            img.save(jpg) 
            # 將已經轉換的圖片移動到指定位置 
            shutil.move(png, output_dir) 
        except IOError as e:
           print('could not read:',jpg) 
           print('error:',e) 
           print('skip it\n') 
    sys.stdout.write('Convert Over!\n') 
    sys.stdout.flush() 
 
def rename(path):
    #給圖片編隊函數
    filelist = os.listdir(path)  #獲取文件路徑
    total_num = len(filelist)  #獲取文件長度
    i = 1  #文件從1開始命名
    for item in filelist:
        if item.endswith('.jpg'): 
            src = os.path.join(os.path.abspath(path), item)
            #dst = os.path.join(os.path.abspath(path), ''+str(i) + '.jpg') 
            dst = os.path.join(os.path.abspath(path), '00' + format(str(i), '0>3s') + '.jpg')   
        try:
            os.rename(src, dst)
            print ('converting %s to %s ...' % (src, dst))
            i = i + 1
        except:
            continue
        print ('total %d to rename & converted %d jpgs' % (total_num, i))
 
    
if __name__ == "__main__": 
  current_dir = os.getcwd() 
  print(current_dir) 
  data_dir = 'data/'  #待轉化圖像文件
  image2png(data_dir,'jpg')
  rename(data_dir)

圖像標註比較消耗人的耐力和專注力，廢話不多話。這次缺陷檢測實驗使用的圖庫爲：DAGM 2007的數據集中的其中三類，圖像標註推薦使用 labelImg軟件，關於這個軟件的安裝和使用參考我的一篇博文：labelImg的安裝和使用，內容很詳細，如果出問題，請留言。標註的過程大概是這樣：

 

3、在github上下載tiny-yolo v3工程

將標註後的圖片和xml文檔分別放到tiny-yolo v3文件下的VOC2007文件下的JPEGImages文件和Annotation文件下。

4、編譯程序運行

使用TensorFlow的編譯器Spyder運行VOC2007文件下的test.py程序，會在VOC2007/ImageSets/Main下生產如下幾個txt文檔，就是對圖片路徑按照比例分成訓練，驗證，測試集，測試集用不上。

 

 

5、

 

對tiny-yolo v3下的voc_annotation.py進行簡單修改，並運行。在tiny-yolo v3下會生成幾個txt文檔，比如：2007_train.txt，這時候我們將前面的2007_刪掉。因爲訓練的缺陷圖像爲三類，voc_annotation.py修改如下：
 

import xml.etree.ElementTree as ET
from os import getcwd
 
sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
 
classes = ["class1", "class2", "class3"]
 
 
def convert_annotation(year, image_id, list_file):
    in_file = open('VOC%s/Annotations/%s.xml'%(year, image_id))
    tree=ET.parse(in_file)
    root = tree.getroot()
 
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
 
wd = getcwd()
 
for year, image_set in sets:
    image_ids = open('VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOC%s/JPEGImages/%s.jpg'%(wd, year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()

 

 

更新補充：

寫的時候忘記了一步，還有一處需要修改。路徑 yolo/model_data下的voc_classes.txt也要修改，因爲我們這次做的三個類別class1,class2,class3的檢測，所以修改如下：

如果要對其他檢測，需要把三個類別名字改成其他的，注意字符的分行，建議使用notepad++修改，因爲python在判定類別的時候使用的是len()函數。

5.1、前面不是說過了嗎，在劃分驗證，測試，訓練集時，訓練和測試集在實際程序運行中根本沒有用上，因爲在程序tiny-train.py中有一個函數已經說明了問題，如下圖，代碼中的annotation_path輸入的是train.txt，然後對其中的圖片重新劃分訓練和測試集。總數×val_split就是訓練圖片數量。

所以我們使用Notepad++打開test.txt和val.txt，將其中的路徑複製到train.txt中，最後只保留train.txt，如下：
圖1 tiny-train.py

 

 

 

複製路徑

 

6、準備訓練

然後對一些代碼文件進行修改，就可以開始訓練了，這裏就不一一敘述具體的修改了。直接在博客前言的github中下載修改好的代碼即可，歡迎加星，有問題的話請提問，會協助解決。

訓練使用：tiny-train.py，批量測試使用：yolo-test-batch.py，將待測試的圖像放在VOC2007文件下的Images文件中。測試結果會在VOC2007下的SegmentationClass文件中。待測圖像使用步驟1進行轉格式和編隊、resize。放一張訓練過程的圖和測試結果圖：

 

 

 

程序是在ubantu下跑的，當然windows系統也可以，只是文件文件需要修改，碼字不易，歡迎交流、點贊，給github加星。

當然，稍微修改也可以測試YOLO v3，和是否加載預權重的程序，以及對其中的特徵網絡等部件修剪，博主都做了。