端口映射 Apache Pivot教程 elasticsearch keras foreach datatables path get textview ionic3 vue例子 vue双向绑定 建站一条龙 bootstrap时间轴 js空格符 ceb转换成pdf在线转换 js教程文档 matlab不等于怎么表示 python爬虫包 cmd清空命令 flutter优缺点 python自学教程 python环境配置 java连接数据库代码 javascript基础 volist 图片链接生成器 js倒计时代码 python队列 pr转场特效下载 saminside 超级力量2修改 矩阵分析与应用 js小数点保留2位 华为下拉开关设置 手机电脑模拟器 python电子书 软件龙头股 3dmax2014下载 office2010免费版
当前位置: 首页 > 学习教程  > 编程语言

python3 TensorRT 推理Demo

2020/12/28 18:48:31 文章标签:

1. 使用tensorrt 对人脸68个特征点推理demo 代码 import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np import cv2def get_engine(engine_path):# If a serialized engine exists, use it instead of building an engine.print(&…

1. 使用tensorrt 对人脸68个特征点推理demo 代码

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2

def get_engine(engine_path):
    # If a serialized engine exists, use it instead of building an engine.
    print("Reading engine from file {}".format(engine_path))
    with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

TRT_LOGGER = trt.Logger()
# engine = get_engine("yolov4_1.trt")


engine = get_engine("pfld.engine")
print(engine)
for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * 1
        dims = engine.get_binding_shape(binding)
        print(size)
        print(dims)
        print(binding)
        print(engine.binding_is_input(binding))
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        print("dtype = ", dtype)



engine = get_engine("pfld.engine")
context = engine.create_execution_context()

def get_landmarks(img):
    resized = cv2.resize(img, (112, 112), interpolation=cv2.INTER_LINEAR)
    img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    cv2.imwrite("tmp.jpg", img_in)
    img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
    img_in = np.expand_dims(img_in, axis=0)
    img_in /= 255.0
    # img_in = np.ascontiguousarray(img_in)
    print("Shape of the network input: ", img_in.shape)
    # print(img_in)

    # with get_engine("mobilefacenet-res2-6-10-2-dim512/onnx/face_reg_mnet.engine") as engine, engine.create_execution_context() as context:
    h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
    h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)

    # Allocate device memory for inputs and outputs.
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    # Create a stream in which to copy inputs/outputs and run inference.
    stream = cuda.Stream()

    # set the host input data
    # h_input = img_in
    np.copyto(h_input, img_in.ravel())
    # np.copyto(h_input, img_in.unsqueeze_(0))

    # print(h_input)
    # Transfer input data to the GPU.
    cuda.memcpy_htod_async(d_input, h_input, stream)
    # Run inference.
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    # Synchronize the stream
    stream.synchronize()
    # Return the host output. 

    # print(h_output)
    return h_output


img1 = cv2.imread("./s_28.jpg")
print(img1.shape)
output = get_landmarks(img1).reshape(-1, 2)

print(output)
for xy in output:
    x = xy[0] * 112
    y = xy[1] * 112
    cv2.circle(img1, (int(x), int(y)), 2, (0,255,0), -1)

cv2.imwrite("out.jpg", img1)

本文链接: http://www.dtmao.cc/news_show_550118.shtml

附件下载

相关教程

    暂无相关的数据...

共有条评论 网友评论

验证码: 看不清楚?