你觉得这篇文章怎么样? 帮助我们为您提供更好的内容。
Thank you! Your feedback has been received.
There was a problem submitting your feedback, please try again later.
你觉得这篇文章怎么样?
作者 |
張嘉鈞 |
難度 |
普通 |
材料表 |
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