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)