python硬解码 opencv gpu解码图片torchnvjpg

torch中有一个包torchnvjpg，可以实现gpu解码图片，解码速度会快于cv2.imread()。

安装

git clone https://github.com/itsliupeng/torchnvjpeg.git

cd torchnvjpeg/

python3 setup.py bdist_wheel

cd dist

python3 -m pip install torchnvjpeg-0.1.0-cp310-cp310-linux_x86_64.whl

使用

import torch
import torchnvjpeg
import torch.nn.functional as F

decoder_gpu = torchnvjpeg.Decoder(device_id=0)

jpgname = 'test.jpg'
image_data = open(jpgname, 'rb').read()

# gpu解码
hwc_image_uint8 = decoder_gpu.decode(image_data)  # gpu数据，uint8类型，颜色通道RGB
src_h,src_w,src_c = hwc_image_uint8.shape
chw_image_float = hwc_image_uint8.permute(2,0,1).float() # hwc 2 chw
bchw_img = chw_image_float.unsqueeze(0) # chw 2 bchw

# resize 将图片resize到224*224
new_size = 224,224
resize_img = F.interpolate(bchw_img, size=new_size,mode='bilinear',align_corners=False)

# pad 在图片周围填充常数114
# 上下左右分别填充10，20，30，40
top, bottom, left, right = 10, 20, 30, 40 
img = F.pad(resize_img, (left,right,top,bottom), 'constant', 114)

# 归一化
img /= 255.0  # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
   img = img.unsqueeze(0)

# 初始化自己的模型
model = ResNet()
model.cuda().eval()

# 输入模型推理
preds = model(img)