Skip to main content
shopping_basket 购物车 0
登录

NVIDIA Jetson Nano應用-使用Pose Estimation進行姿態辨識以及手部辨識

CAVE Education
0

 

作者

張嘉鈞

難度

普通

材料表

Webcam X1

NVIDIA JetsonNano X1

 

 

首先,我們先來運行 trt_pose,這個專案是一個即時姿態辨識的範例!他除了有提供Pre-Trained Model外也有提供Jupyter Notebook形式的教學程式,除此之外他使用的數據集是MSCOCO,所以也有大家也可以嘗試去訓練自己的數據,只需要參照MSCOCO的格式走即可,接下來就教大家怎麼部屬在Jetson Nano上吧。

 

注意:此篇文章使用的Jetson Nano系統為 JetPack 4.4.1。

 

1.安裝相依套件

PyTorch、torchvision:

首先最主要的就是PyTorch、torchvision這邊我們可以參考原廠的作法,我提供的範例為 pytorch 1.6,記得要確認自己的JetPack版本與PyTorch的對應版本是否正確。

## PyTorch
wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
pip3 install Cython
pip3 install numpy torch-1.6.0-cp36-cp36m-linux_aarch64.whl

## torchvision
$ sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
$ git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision   # see below for version of torchvision to download
$ cd torchvision
$ export BUILD_VERSION=0.7.0  # where 0.x.0 is the torchvision version  
$ python3 setup.py install --user

 

 

torch2trt:

NVIDIA的專案基本上都會轉換成TensorRT引擎,讓整個推論的速度加快,而我們這次使用的trt_pose他提供的模型是PyTorch的,所以我們可以直接用這個Github提供的程式來做轉換,透過torch2trt的函式庫我們也將能更方便的使用TensorRT引擎,使用方法會與PyTorch相似。

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
sudo python3 setup.py install --plugins

 

其他相關套件(tqdm, cython, pycocotools, matplotlib):

進度條使用到tqdm;cython則是Python跟C/C++的橋樑;由於我們使用到coco dataset所以也可以安裝Python的工具,最後Jupyter Lab與matplotlib相容度最高,所以在上面開圖片基本上都會使用該函式庫。

sudo pip3 install tqdm cython pycocotools
sudo apt-get install python3-matplotlib

 

2. 下載安裝 trt_pose的Github:

cd ~
git clone https://github.com/NVIDIA-AI-IOT/trt_pose
cd trt_pose
sudo python3 setup.py install

 

3.下載預訓練模型:

我們需要將模型下載到 ~/trt_pose/tasks/human_pose當中,由於Github提供的連結是Google Drvie不能直接透過wget加上網址去下載,所以我會建議直接從Github上下載,或者使用下列指令去下載,這個指令可以下載低於100M的Google Drive檔。

wget --no-check-certificate "https://drive.google.com/uc?export=download&id=1XYDdCUdiF2xxx4rznmLb62SdOUZuoNbd" -O ./tasks/human_pose/resnet18_baseline_att_224x224_A.pth

wget --no-check-certificate "https://drive.google.com/uc?export=download&id=13FkJkx7evQ1WwP54UmdiDXWyFMY1OxDU" -O ./tasks/human_pose/densenet121_baseline_att_256x256_B.pth

 

4.運行程式:

原廠有提供 trt_pose/tasks/human_pose/live_demo.ipynb,不過我決定帶大家寫兩支程式,順便熟悉各個指令的用途,首先我們需要將下載下來的PyTorch模型轉換成TensorRT引擎,所以我們先新增一個檔案名為 cvt_trt.py,接下來就一一新增程式碼,首先導入函式庫:

import json

import torch2trt
import torch

import trt_pose.coco
import trt_pose.models

def print_div(txt):    
    print(f"\n{txt}\n")

 

接著載入預訓練模型並且修改最後一層,這邊增加的是 cmap_channels跟 paf_channels,每一個模型會有兩個輸出cmap與paf,cmap(Confidence Map)就是Pose Estimation常見的熱力圖,能透過其找出Keypoint所在;Paf (Part Affinity Fields)則是提供一個向量區域,可以表示人體各個Keypoint、skeleton的關聯性,而paf 由於是一個向量所以會有 x, y的數值,輸出要乘以2。

print_div("LOAD MODEL")

with open('human_pose.json', 'r') as f:
    human_pose = json.load(f)

# 取得 keypoint 數量
num_parts = len(human_pose['keypoints'])
num_links = len(human_pose['skeleton'])

# 修改輸出層
model = trt_pose.models.resnet18_baseline_att(num_parts, 2 * num_links).cuda().eval()

# 載入權重
MODEL_WEIGHTS = 'resnet18_baseline_att_224x224_A.pth'
model.load_state_dict(torch.load(MODEL_WEIGHTS))

 

關於cmap與paf的詳細可以去看這篇論文的敘述Realtime Multi-Person 2D Pose Estimation using Part Affinity Fields。

 

圖片擷取自Realtime Multi-Person 2D Pose Estimation using Part Affinity Fields 論文

 

最後我們就可以進行轉換與儲存引擎:

print_div("COVERTING")

WIDTH, HEIGHT = 224, 224
data = torch.zeros((1, 3, HEIGHT, WIDTH)).cuda()
model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)

print_div("SAVING TENSORRT")

OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A.trt'
torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)

print_div("FINISH")

 

儲存檔案之後就可以透過Python3去運行程式來將PyTorch轉換成TensorRT:

$ python3 cvt_trt.py

 

我們已經完成第一步轉換模型了,接下來要再新增一個Python檔案 demo.py,來實作推論的部分,首先一樣先導入函式庫:

import torch
import torchvision.transforms as transforms

import cv2
import PIL.Image
import json
import time

from torch2trt import TRTModule
import trt_pose
import trt_pose.coco
import trt_pose.models
from trt_pose.draw_objects import DrawObjects
from trt_pose.parse_objects import ParseObjects

 

接著我定義了兩個副函式,一個是print_div與cvt_trt.py中功能相同只是換了寫法;另一個是preprocess用來將OpenCV的圖片轉換成PIL並且進行正規化:

print_div = lambda x: print(f"\n{x}\n")

mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda()
device = torch.device('cuda')

def preprocess(image):
    global device
    device = torch.device('cuda')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device)
    image.sub_(mean[:, None, None]).div_(std[:, None, None])
    return image[None, ...]

 

接著載入TensorRT引擎,通過torch2trt的函式庫可以使用類似PyTorch的方式匯入,更簡單直觀:

print_div("LOAD TENSORRT ENGINE")

OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A.trt'

model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))

 

接著就到Real-time的Inference程式,我們透過OpenCV擷取攝影機,並且先取得解析度,待會會將圖片進行裁切、縮放:

print_div("START STREAM")

cap = cv2.VideoCapture(0)

w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
crop_size = (w-h)//2
# print(w, h, crop_size)

 

我們需要透過While迴圈讓攝影機不斷擷取當前的frame並且進行Inference以及顯示結果,t_start用於計算FPS,如果沒有影像的話我們就會跳過當前迴圈,直到有影像內容才會往下運行:

while(True):
    
    t_start = time.time()
    
    ret, frame = cap.read()
    
    if not ret:
        continue

 

資料前處理的部分除了剛剛撰寫的preprocess之外還需要針對frame裁切成正方形以及縮放大小,若省略這些步驟會導致模型辨識不出結果:

    frame = frame[:, crop_size:(w-crop_size)]
    image = cv2.resize(frame, (224,224))
    data = preprocess(image)

 

接著就是推論的部分,模型丟入資料後會輸出cmap以及paf,透過 parse_objects可以解析出裡面的內容,counts是物件的數量;objects是物件的座標等資訊;peaks用於繪製skeleton。

    cmap, paf = model_trt(data)
    cmap, paf = cmap.detach().cpu(), paf.detach().cpu()
    counts, objects, peaks = parse_objects(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15)
    draw_objects(image, counts, objects, peaks)

 

最後就是顯示結果以及一些基本的關閉設定,這裡我們會運算出FPS並將其繪製在圖片上:

    t_end = time.time()
    cv2.putText(image, f"FPS:{int(1/(t_end-t_start))}", (10,20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1,  cv2.LINE_AA)
    cv2.imshow('pose esimation', image)

    if cv2.waitKey(1)==ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

 

執行結果如下,可以看到Jetson Nano運行有12左右的FPS,效果還算不錯:

 

到目前我們已經成功執行了trt_pose,而強大的作者們還有針對手部做一個專案,我們就直接來安裝並執行看看吧。

 

1.安裝 trt_pose以及相依套件,如果從文章開始執行到現在基本上這一步都處理好了,有需求的讀者可以往回查看。

 

2.安裝traitlets:

$ pip3 install traitlets

 

3.下載他們的Github並移動到該目錄:

$ cd ~
$ git clone https://github.com/NVIDIA-AI-IOT/trt_pose_hand.git && cd trt_pose_hand

 

4.下載預訓練模型:

$ wget --no-check-certificate "https://drive.google.com/uc?export=download&id=1NCVo0FiooWccDzY7hCc5MAKaoUpts3mo" -O ./model /hand_pose_resnet18_baseline_att_224x224.pth

 

接著就可以撰寫程式來運行了,他們一樣有提供範例程式live_hand_demo.ipynb,不過我們也一樣自己來寫一個吧!首先也需要將PyTorch的預訓練模型給轉換好,新增一個cvt_trt.py,程式碼與早先的那個幾乎一樣,只需要更改json、.pth、.trt的名稱即可:

import json

import torch2trt
import torch

import trt_pose.coco
import trt_pose.models

def print_div(txt):    
    print(f"\n{txt}\n")

print_div("LOAD MODEL")

with open('preprocess/hand_pose.json', 'r') as f:
    human_pose = json.load(f)

# 取得 keypoint 數量
num_parts = len(human_pose['keypoints'])
num_links = len(human_pose['skeleton'])

# 修改輸出層
model = trt_pose.models.resnet18_baseline_att(num_parts, 2 * num_links).cuda().eval()

# 載入權重
MODEL_WEIGHTS = 'model/hand_pose_resnet18_baseline_att_224x224.pth'
model.load_state_dict(torch.load(MODEL_WEIGHTS))

print_div("COVERTING")

WIDTH, HEIGHT = 224, 224
data = torch.zeros((1, 3, HEIGHT, WIDTH)).cuda()
model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)

print_div("SAVING TENSORRT")

OPTIMIZED_MODEL = 'model/hand_pose_resnet18_baseline_att_224x224.trt'
torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)

print_div("FINISH")

 

接著新增demo.py,也與之前的雷同,json、trt的檔名記得改,值得注意的地方是Parse Objects 這個函式多設定了cmap_threshold與 link_threshold,因為這個模型其他地方很容易辨識錯誤,所以只好提高閥值:

import torch
import torchvision.transforms as transforms

import cv2
import PIL.Image
import json
import time

import math
import os
import numpy as np
import traitlets

from torch2trt import TRTModule
import trt_pose
import trt_pose.coco
import trt_pose.models
from trt_pose.draw_objects import DrawObjects
from trt_pose.parse_objects import ParseObjects

#############################################################

print_div = lambda x: print(f"\n{x}\n")

mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda()
device = torch.device('cuda')

def preprocess(image):
    global device
    device = torch.device('cuda')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device)
    image.sub_(mean[:, None, None]).div_(std[:, None, None])
    return image[None, ...]

############################################################

print_div("INIT")

with open('preprocess/hand_pose.json', 'r') as f:
    hand_pose = json.load(f)

# 改變標籤結構 : 增加頸部keypoint以及 paf的維度空間
topology = trt_pose.coco.coco_category_to_topology(hand_pose)
# 用於解析預測後的 cmap與paf
parse_objects = ParseObjects(topology, cmap_threshold=0.30, link_threshold=0.30)
# 用於將keypoint繪製到圖片上
draw_objects = DrawObjects(topology)

############################################################

print_div("LOAD TENSORRT ENGINE")

OPTIMIZED_MODEL = 'model/hand_pose_resnet18_baseline_att_224x224.trt'

model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))

###########################################################

print_div("START STREAM")

cap = cv2.VideoCapture(0)

w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
crop_size = (w-h)//2
# print(w, h, crop_size)

while(True):
    
    t_start = time.time()
    
    ret, frame = cap.read()
    
    if not ret:
        continue

    frame = frame[:, crop_size:(w-crop_size)]
    image = cv2.resize(frame, (224,224))
    data = preprocess(image)

    cmap, paf = model_trt(data)
    cmap, paf = cmap.detach().cpu(), paf.detach().cpu()
    counts, objects, peaks = parse_objects(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15)
    draw_objects(image, counts, objects, peaks)

    t_end = time.time()
    cv2.putText(image, f"FPS:{int(1/(t_end-t_start))}", (10,20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1,  cv2.LINE_AA)
    cv2.imshow('pose esimation', image)

    if cv2.waitKey(1)==ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

#############################################################

 

 

執行結果如下:

 

結語

恭喜你們成功執行trt_pose以及trt_pose_hand,針對trt_pose_hand作者還有兩個範例可以運行,我們將在下一篇進行介紹與分析。

 

相關文章

trt_pose

https://github.com/NVIDIA-AI-IOT/trt_pose

 

Hand Pose Estimation And Classification

https://github.com/NVIDIA-AI-IOT/trt_pose_hand

CAVEDU Education is devoted into robotics education and maker movement since 2008, and is intensively active in teaching fundamental knowledge and skills. We had published many books for readers in all ages, topics including Deep Learning, edge computing, App Inventor, IoT and robotics. Please check CAVEDU's website for more information: http://www.cavedu.com, http://www.appinventor.tw

相关内容

DesignSpark Electrical Logolinkedin