helloproject-ai/test_script/nvjpeg_infer.py

33 lines
1.2 KiB
Python

import time
import tkinter
# import cv2
from nvjpeg_decoder import decode
from os import listdir, getcwd
from os.path import join
from numpy import array, fromfile, uint8
# from matplotlib import pyplot, figure
# from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
# import matplotlib_fontja
from onnxruntime import InferenceSession
onnx_path = r"C:\Users\tomokazu\RustroverProjects\ameba_blog_downloader\src\retinaface\resnet_retinaface.onnx"
datadir = r"D:\helloproject-ai-data\blog_images"
session = InferenceSession(
path_or_bytes=onnx_path,
# providers=[
# # 'TensorrtExecutionProvider',
# # 'CUDAExecutionProvider',
# 'CPUExecutionProvider'
# ]
)
for member in listdir(datadir):
for file in listdir(join(datadir, member)):
# with open(join(datadir, member, file), mode="rb") as f:
(data, (scale, (width, height))) = decode(fromfile(join(datadir, member, file), dtype=uint8), "imagenet",
(1080, 1080))
print(width, height)
image_arr = array(data).reshape((1, 3, height, width)) # .transpose([1, 2, 0])
session.run(input_feed={'input': image_arr}, output_names=['bbox', 'confidence', 'landmark'])