如何在 MediaPipe 中使用串流影像產生高品質的辨識結果

想要使用 MediaPipe 產生高品質的輸出結果,必須使用 LIVE_STREAM 參數也就是串流影像,但官方文件只介紹到這裡就沒有進一步說明如何撰寫。為方便有需要的人快速上手,這裡以手部地標為範例,說明如何使用 LIVE_STREAM 進行辨識,其他主題,原則上舉一反三即可。

先載入需要的函數庫,然後宣告兩個全域變數 result_queue 與 HandLandmarkerResult。前者用來讓 OpenCV 強制等待 MediaPipe 辨識完成才能顯示結果,也就是讓 MediaPipe 的非同步辨識變成同步;後者則是 MediaPipe 的非同步函數中第一個參數需要的資料型態,稍後會看到這個函數。

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2
import time
import queue

result_queue = queue.Queue(1)
HandLandmarkerResult = mp.tasks.vision.HandLandmarkerResult

接下來撰寫一個 MediaPipe 手部地標偵測的初始化函數,這邊的程式碼基本上複製貼上官方文件中的程式碼,只是其中的 HandLandmarkerResult 變數被移到全域位置。在 options 參數中必須使用 LIVE_STREAM,這樣 MediaPipe 才會啟動移動追蹤功能,也就是隔幾個畫面才會真正用類神經網路計算一次手部地標位置,其餘都是移動追蹤,這樣才能產生高品質的輸出結果。參數中的 result_callback 用來註冊一個回呼函數 print_result() 負責處理計算結果。手部地標辨識的模型檔(.task)可在官網中找到,請先下載回來。

def init():    
    BaseOptions = mp.tasks.BaseOptions
    HandLandmarker = mp.tasks.vision.HandLandmarker
    HandLandmarkerOptions = mp.tasks.vision.HandLandmarkerOptions
    VisionRunningMode = mp.tasks.vision.RunningMode

    options = HandLandmarkerOptions(
        base_options=BaseOptions(model_asset_path='hand_landmarker.task'),
        running_mode=VisionRunningMode.LIVE_STREAM,
        num_hands=2,
        result_callback=print_result
    )
    landmarker = HandLandmarker.create_from_options(options)
    return landmarker

接著定義 print_result() 函數。其中第三行要將 MediaPipe 的 Image 物件透過 numpy_view() 函數轉成 numpy 格式,也就是 OpenCV 的格式,最後加上一個 copy() 函數,因為 numpy_view() 轉出來的資料是 read-only,後續無法在轉出來的圖形上畫上需要的圖案。兩個 for 迴圈中將手指的每個關節座標畫出,並且分左右手不同顏色,然後食指的指尖用紅色特別標出,最後將處理好的畫面放進佇列中。迴圈內的程式碼完全根據 result 結果來設計,可將第二行的註解打開,仔細觀察 result 的輸出內容。

def print_result(result: HandLandmarkerResult, output_image: mp.Image, timestamp_ms: int):
   # print('hand landmarker result: {}'.format(result))
   frame = output_image.numpy_view().copy()
   for i, landmark in enumerate(result.hand_landmarks):
       hand_name = result.handedness[i][0].category_name
       for j, point in enumerate(landmark):
           x = int(point.x * frame.shape[1])
           y = int(point.y * frame.shape[0])
           if hand_name == 'Left':
               cv2.circle(frame, (x, y), 3, (255,255,0), -1)
           else:
               cv2.circle(frame, (x, y), 3, (0,255,255), -1)
           if j == 8:   # 食指
               cv2.circle(frame, (x, y), 10, (0,0,255), -1)
   result_queue.put(frame)

接下來主程式了,這裡的程式碼基本上就是 OpenCV 的攝影機影像輸入,然後轉交給 MediaPipe 去分析與辨識,最後將結果從佇列 result_queue 中取出並顯示,如果佇列目前是空的,get() 函數會進入等待。

def main(landmarker):
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    while True:
        begin_time = time.time()
        ret, frame = cap.read()
        frame = cv2.flip(frame, 1)

        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
        landmarker.detect_async(mp_image, int(time.time() * 1000))

        frame = result_queue.get()
        cv2.imshow('frame', frame)
        print(f'{1 / (time.time() - begin_time):.2f}')

        if cv2.waitKey(1) == 27:
            break

    cv2.destroyAllWindows()

最後啟動系統。

landmarker = init()
main(landmarker)

執行看看,應該會得到一個非常好的執行結果,但前提是 CPU 不能太差,若 FPS 不到 5,還是會有明顯的延遲感,但不論畫面是否延遲,手部地標的辨識結果依舊非常穩定。

發表迴響