Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
GoAI
handpose_x
提交
cf5217ad
handpose_x
项目概览
GoAI
/
handpose_x
与 Fork 源项目一致
Fork自
DataBall / handpose_x
通知
1
Star
0
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
handpose_x
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
cf5217ad
编写于
10月 13, 2021
作者:
L
lixiang
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
add onnx tool
上级
3f071225
变更
3
隐藏空白更改
内联
并排
Showing
3 changed file
with
227 addition
and
2 deletion
+227
-2
README.md
README.md
+4
-2
model2onnx.py
model2onnx.py
+110
-0
onnx_inference.py
onnx_inference.py
+113
-0
未找到文件。
README.md
浏览文件 @
cf5217ad
...
...
@@ -2,8 +2,7 @@
手势 21 个关键点检测
## 重要更新
### 第一个完整pipeline 的 "dpcas" 项目,本地手势交互应用 Demo,之后会推出web架构的手势交互。
### 该项目 "dpcas" 虽然它才刚刚诞生,有各种不足,但是我会继续改进,努力让更多人看到它,希望它不仅仅是一个demo。
### 添加 onnx 模块,预训练模型中有转好的resnet50-onnx模型,注意:目前不支持rexnetv1
### "dpcas" 项目地址:https://codechina.csdn.net/EricLee/dpcas
## 项目Wiki
...
...
@@ -104,6 +103,9 @@
### 模型推理
*
根目录下运行命令: python inference.py (注意脚本内相关参数配置 )
### onnx使用
*
step1: 设定相关配置包括模型类型和模型参数路径,根目录下运行命令: python model2onnx.py (注意脚本内相关参数配置 )
*
step2: 设定onnx模型路径,根目录下运行命令: python onnx_inference.py (注意脚本内相关参数配置 )
*
建议
```
...
...
model2onnx.py
0 → 100644
浏览文件 @
cf5217ad
#-*-coding:utf-8-*-
# date:2021-10-5
# Author: Eric.Lee
# function: pytorch model 2 onnx
import
os
import
argparse
import
torch
import
torch.nn
as
nn
import
numpy
as
np
from
models.resnet
import
resnet18
,
resnet34
,
resnet50
,
resnet101
from
models.squeezenet
import
squeezenet1_1
,
squeezenet1_0
from
models.shufflenetv2
import
ShuffleNetV2
from
models.shufflenet
import
ShuffleNet
from
models.mobilenetv2
import
MobileNetV2
from
torchvision.models
import
shufflenet_v2_x1_5
,
shufflenet_v2_x1_0
,
shufflenet_v2_x2_0
from
models.rexnetv1
import
ReXNetV1
if
__name__
==
"__main__"
:
parser
=
argparse
.
ArgumentParser
(
description
=
' Project Hand Pose Inference'
)
parser
.
add_argument
(
'--model_path'
,
type
=
str
,
default
=
'./weights1/resnet_50-size-256-wingloss102-0.119.pth'
,
help
=
'model_path'
)
# 模型路径
parser
.
add_argument
(
'--model'
,
type
=
str
,
default
=
'shufflenet_v2_x1_5'
,
help
=
'''model : resnet_34,resnet_50,resnet_101,squeezenet1_0,squeezenet1_1,shufflenetv2,shufflenet,mobilenetv2
shufflenet_v2_x1_5 ,shufflenet_v2_x1_0 , shufflenet_v2_x2_0'''
)
# 模型类型
parser
.
add_argument
(
'--num_classes'
,
type
=
int
,
default
=
42
,
help
=
'num_classes'
)
# 手部21关键点, (x,y)*2 = 42
parser
.
add_argument
(
'--GPUS'
,
type
=
str
,
default
=
'0'
,
help
=
'GPUS'
)
# GPU选择
parser
.
add_argument
(
'--test_path'
,
type
=
str
,
default
=
'./image/'
,
help
=
'test_path'
)
# 测试图片路径
parser
.
add_argument
(
'--img_size'
,
type
=
tuple
,
default
=
(
256
,
256
),
help
=
'img_size'
)
# 输入模型图片尺寸
parser
.
add_argument
(
'--vis'
,
type
=
bool
,
default
=
True
,
help
=
'vis'
)
# 是否可视化图片
print
(
'
\n
/******************* {} ******************/
\n
'
.
format
(
parser
.
description
))
#--------------------------------------------------------------------------
ops
=
parser
.
parse_args
()
# 解析添加参数
#--------------------------------------------------------------------------
print
(
'----------------------------------'
)
unparsed
=
vars
(
ops
)
# parse_args()方法的返回值为namespace,用vars()内建函数化为字典
for
key
in
unparsed
.
keys
():
print
(
'{} : {}'
.
format
(
key
,
unparsed
[
key
]))
#---------------------------------------------------------------------------
os
.
environ
[
'CUDA_VISIBLE_DEVICES'
]
=
ops
.
GPUS
test_path
=
ops
.
test_path
# 测试图片文件夹路径
#---------------------------------------------------------------- 构建模型
print
(
'use model : %s'
%
(
ops
.
model
))
if
ops
.
model
==
'resnet_50'
:
model_
=
resnet50
(
num_classes
=
ops
.
num_classes
,
img_size
=
ops
.
img_size
[
0
])
elif
ops
.
model
==
'resnet_18'
:
model_
=
resnet18
(
num_classes
=
ops
.
num_classes
,
img_size
=
ops
.
img_size
[
0
])
elif
ops
.
model
==
'resnet_34'
:
model_
=
resnet34
(
num_classes
=
ops
.
num_classes
,
img_size
=
ops
.
img_size
[
0
])
elif
ops
.
model
==
'resnet_101'
:
model_
=
resnet101
(
num_classes
=
ops
.
num_classes
,
img_size
=
ops
.
img_size
[
0
])
elif
ops
.
model
==
"squeezenet1_0"
:
model_
=
squeezenet1_0
(
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"squeezenet1_1"
:
model_
=
squeezenet1_1
(
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"shufflenetv2"
:
model_
=
ShuffleNetV2
(
ratio
=
1.
,
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"shufflenet_v2_x1_5"
:
model_
=
shufflenet_v2_x1_5
(
pretrained
=
False
,
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"shufflenet_v2_x1_0"
:
model_
=
shufflenet_v2_x1_0
(
pretrained
=
False
,
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"shufflenet_v2_x2_0"
:
model_
=
shufflenet_v2_x2_0
(
pretrained
=
False
,
num_classes
=
ops
.
num_classes
)
elif
ops
.
model
==
"shufflenet"
:
model_
=
ShuffleNet
(
num_blocks
=
[
2
,
4
,
2
],
num_classes
=
ops
.
num_classes
,
groups
=
3
)
elif
ops
.
model
==
"mobilenetv2"
:
model_
=
MobileNetV2
(
num_classes
=
ops
.
num_classes
)
use_cuda
=
torch
.
cuda
.
is_available
()
device
=
torch
.
device
(
"cuda:0"
if
use_cuda
else
"cpu"
)
model_
=
model_
.
to
(
device
)
model_
.
eval
()
# 设置为前向推断模式
# 加载测试模型
if
os
.
access
(
ops
.
model_path
,
os
.
F_OK
):
# checkpoint
chkpt
=
torch
.
load
(
ops
.
model_path
,
map_location
=
device
)
model_
.
load_state_dict
(
chkpt
)
print
(
'load test model : {}'
.
format
(
ops
.
model_path
))
input_size
=
ops
.
img_size
[
0
]
batch_size
=
1
#批处理大小
input_shape
=
(
3
,
input_size
,
input_size
)
#输入数据,改成自己的输入shape
print
(
"input_size : "
,
input_size
)
x
=
torch
.
randn
(
batch_size
,
*
input_shape
)
# 生成张量
x
=
x
.
to
(
device
)
export_onnx_file
=
"{}_size-{}.onnx"
.
format
(
ops
.
model
,
input_size
)
# 目的ONNX文件名
torch
.
onnx
.
export
(
model_
,
x
,
export_onnx_file
,
opset_version
=
9
,
do_constant_folding
=
True
,
# 是否执行常量折叠优化
input_names
=
[
"input"
],
# 输入名
output_names
=
[
"output"
],
# 输出名
#dynamic_axes={"input":{0:"batch_size"}, # 批处理变量
# "output":{0:"batch_size"}}
)
onnx_inference.py
0 → 100644
浏览文件 @
cf5217ad
#-*-coding:utf-8-*-
# date:2021-10-5
# Author: Eric.Lee
# function: onnx Inference
import
os
,
sys
sys
.
path
.
append
(
os
.
getcwd
())
import
onnxruntime
import
onnx
import
cv2
import
torch
import
numpy
as
np
from
hand_data_iter.datasets
import
draw_bd_handpose
class
ONNXModel
():
def
__init__
(
self
,
onnx_path
,
gpu_cfg
=
False
):
"""
:param onnx_path:
"""
self
.
onnx_session
=
onnxruntime
.
InferenceSession
(
onnx_path
)
if
gpu_cfg
:
self
.
onnx_session
.
set_providers
([
'CUDAExecutionProvider'
],
[
{
'device_id'
:
0
}])
self
.
input_name
=
self
.
get_input_name
(
self
.
onnx_session
)
self
.
output_name
=
self
.
get_output_name
(
self
.
onnx_session
)
print
(
"input_name:{}"
.
format
(
self
.
input_name
))
print
(
"output_name:{}"
.
format
(
self
.
output_name
))
def
get_output_name
(
self
,
onnx_session
):
"""
output_name = onnx_session.get_outputs()[0].name
:param onnx_session:
:return:
"""
output_name
=
[]
for
node
in
onnx_session
.
get_outputs
():
output_name
.
append
(
node
.
name
)
return
output_name
def
get_input_name
(
self
,
onnx_session
):
"""
:param onnx_session:
:return:
"""
input_name
=
[]
for
node
in
onnx_session
.
get_inputs
():
input_name
.
append
(
node
.
name
)
return
input_name
def
get_input_feed
(
self
,
input_name
,
image_numpy
):
"""
:param input_name:
:param image_numpy:
:return:
"""
input_feed
=
{}
for
name
in
input_name
:
input_feed
[
name
]
=
image_numpy
return
input_feed
def
forward
(
self
,
image_numpy
):
'''
# image_numpy = image_numpy[np.newaxis, :]
# onnx_session.run([output_name], {input_name: x})
# :param image_numpy:
# :return:
'''
input_feed
=
self
.
get_input_feed
(
self
.
input_name
,
image_numpy
)
output
=
self
.
onnx_session
.
run
(
self
.
output_name
,
input_feed
=
input_feed
)
return
output
if
__name__
==
"__main__"
:
img_size
=
256
model
=
ONNXModel
(
"resnet_50_size-256.onnx"
)
path_
=
"./image/"
for
f_
in
os
.
listdir
(
path_
):
img0
=
cv2
.
imread
(
path_
+
f_
)
img_width
=
img0
.
shape
[
1
]
img_height
=
img0
.
shape
[
0
]
img
=
cv2
.
resize
(
img0
,
(
img_size
,
img_size
),
interpolation
=
cv2
.
INTER_CUBIC
)
img_ndarray
=
img
.
transpose
((
2
,
0
,
1
))
img_ndarray
=
img_ndarray
/
255.
img_ndarray
=
np
.
expand_dims
(
img_ndarray
,
0
)
output
=
model
.
forward
(
img_ndarray
.
astype
(
'float32'
))[
0
][
0
]
output
=
np
.
array
(
output
)
print
(
output
.
shape
[
0
])
pts_hand
=
{}
#构建关键点连线可视化结构
for
i
in
range
(
int
(
output
.
shape
[
0
]
/
2
)):
x
=
(
output
[
i
*
2
+
0
]
*
float
(
img_width
))
y
=
(
output
[
i
*
2
+
1
]
*
float
(
img_height
))
pts_hand
[
str
(
i
)]
=
{}
pts_hand
[
str
(
i
)]
=
{
"x"
:
x
,
"y"
:
y
,
}
draw_bd_handpose
(
img0
,
pts_hand
,
0
,
0
)
# 绘制关键点连线
#------------- 绘制关键点
for
i
in
range
(
int
(
output
.
shape
[
0
]
/
2
)):
x
=
(
output
[
i
*
2
+
0
]
*
float
(
img_width
))
y
=
(
output
[
i
*
2
+
1
]
*
float
(
img_height
))
cv2
.
circle
(
img0
,
(
int
(
x
),
int
(
y
)),
3
,
(
255
,
50
,
60
),
-
1
)
cv2
.
circle
(
img0
,
(
int
(
x
),
int
(
y
)),
1
,
(
255
,
150
,
180
),
-
1
)
cv2
.
namedWindow
(
'image'
,
0
)
cv2
.
imshow
(
'image'
,
img0
)
if
cv2
.
waitKey
(
600
)
==
27
:
break
cv2
.
waitKey
(
0
)
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录