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