Tensorflow版的YOLOv4轉換教學

YOLOv4是目前非常受到歡迎的物件偵測工具,他本身後端使用的類神經框架是Darknet(https://github.com/pjreddie/darknet),而今天則是要把他轉換成Tensorflow框架,這樣做有什麼好處呢?

1. 一般我們在學習AI演算法的過程,多數已TF框架為入門,而Darknet一般只用來實做YOLO辨識,學生在學習上很難從TF很忽然跳到Darknet上,而轉換後TF語法都可以使用,學習上沒有斷層。

2. Darknet在Windows上編譯GPU版本太過瑣碎,而直接使用OpenCV的DNN模型則無法作到GPU加速,但轉換成TF後,就可以直接使用原本TF GPU加速的設定。經過測試,在1024×768的解析度下,OpenCV(非Cuda加速版)FPS=0.2X,而TF版的FPS可達到2-5,差異大概十倍。

我們使用的轉換工具是tf-yolov4(https://github.com/sicara/tf2-yolov4),根據官方說明,目前沒有Train功能,只能做Inference,所以訓練還是得到Colab(誤)。

以下將過程簡要的說明

1. 使用Anaconda建立Python3.9的虛擬環境:由於轉換過程提示建議使用TF2.6以上,因此以Python3.9 版本來建議虛擬環境,因為依據經驗TF-GPU 2.6 搭配的是Py39,所以才需要建立Python3.9的虛擬環境。若您不使用GPU加速,則不限制Python3.9。

2. 開啟命令視窗

    2.1 安裝Tensorflow-gpu:conda install tensorflow-gpu==2.6

    2.2 安裝轉換工具:pip install tf2-yolov4

3. 下載YOLOv4權重檔:https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights,並放置到資料夾內。

4. 執行轉換命令,請注意yolov4.weights的路徑是否正確:convert-darknet-weights yolov4.weights -o yolov4.h5

若要轉換成其他種類格式,例如tflite請參閱作者的github說明

5. 完成後就可以得到yolov4的TF的H5權重檔:yolov4.h5

接下來就可以直接使用我寫好的程式來執行

import tensorflow as tf
from tf2_yolov4.anchors import YOLOV4_ANCHORS #pip install tf2-yolov4
from tf2_yolov4.model import YOLOv4
import time
import cv2 #建議使用4.5.5版本
print(cv2.__version__)
# 選擇攝影機
# target="https://thbcctv01.thb.gov.tw/T2D-11K+190"
cap = cv2.VideoCapture(0) #cap = cv2.VideoCapture(target) 可以觀察路況

WIDTH, HEIGHT = (640, 480) #選擇辨識畫面解析度

model = YOLOv4(
    input_shape=(HEIGHT, WIDTH, 3), #輸入影像規格
    anchors=YOLOV4_ANCHORS,#使用YOLO設定的錨
    num_classes=80, #辨識物件80種    
    yolo_max_boxes=50, #最多找到50個
    yolo_iou_threshold=0.5, #iou門檻0.5
    yolo_score_threshold=0.5, #信任門檻0.5
)
 
model.load_weights('yolov4.h5') #請注意路徑是否正確
#YOLOv4所能辨識的物件列表 
CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
    'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
    'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
    'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
    'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli',
    'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant',
    'bed', 'dining table', 'toilet', 'tv', 'laptop',  'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

while cap.isOpened():#鏡頭能開啟嗎?
    stime=time.time()
    ret, frame = cap.read()
    frame = cv2.resize(frame,(WIDTH, HEIGHT))
    
    #將cv2影像轉換成TF格式
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image = tf.expand_dims(tf.convert_to_tensor(image, dtype=tf.float32) , axis=0) / 255


    #進行物件偵測
    boxes, scores, classes, null = model.predict(image)
    #boxes物件在畫面的位置
    boxes = boxes[0] * [WIDTH, HEIGHT, WIDTH, HEIGHT]
    #scores物件的信賴程度
    scores = scores[0]
    #物件的名稱
    classes = classes[0].astype(int)

    #依序讀取偵測到的物件,並畫出框線
    for (xmin, ymin, xmax, ymax), score, class_idx in zip(boxes, scores, classes):
        if score > 0.5: #設定信任度>0.5才會顯示
            #畫出物件範圍
            cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (255, 178, 50), 3)
            #在物件範圍左上寫出物件名稱+信任度
            text = CLASSES[class_idx] + ': {0:.2f}'.format(score)            
            cv2.putText(frame, text,  (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 2)
    #計算fps
    etime=time.time()
    fps=round(1/(etime-stime),2)
    #將fps放到圖片左上角
    cv2.putText(frame, str(fps),  (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 2)
    #顯示辨識結果
    cv2.imshow("YOLO", frame)   
    key=cv2.waitKey(1)#0=強制等待、1:等候1ms就跳過  更新影像
    if key & 0xFF == ord('q'): #使用者按了鍵盤'q'
        break

# 釋放攝影機
cap.release()

# 關閉所有 OpenCV 視窗
cv2.destroyAllWindows()

這裡做一個簡單的demo,上圖為使用opencv直接載入yolov4,而下圖則為使用tf-gpu的yolov4,可以發現偵測效果基本沒差異,但是fps則差了快10倍。

本文主要參考:https://lindevs.com/yolov4-object-detection-using-tensorflow-2/

Leave a Comment

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *