1 Star 0 Fork 0

minwin/Yolov8_tensorrt

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
infer-det-without-torch-unchanged.py 2.97 KB
一键复制 编辑 原始数据 按行查看 历史
wangjiawei 提交于 2023-12-19 15:55 . initial v8 infer with trt
import argparse
from pathlib import Path
import cv2
import numpy as np
from config import CLASSES, COLORS
from models.utils import blob, det_postprocess, letterbox, path_to_list
from models.utils import det_postprocess_unchanged_onnx
def main(args: argparse.Namespace) -> None:
if args.method == 'cudart':
from models.cudart_api import TRTEngine
elif args.method == 'pycuda':
from models.pycuda_api import TRTEngine
else:
raise NotImplementedError
Engine = TRTEngine(args.engine)
H, W = Engine.inp_info[0].shape[-2:]
images = path_to_list(args.imgs)
save_path = Path(args.out_dir)
if not args.show and not save_path.exists():
save_path.mkdir(parents=True, exist_ok=True)
for image in images:
save_image = save_path / image.name
bgr = cv2.imread(str(image))
draw = bgr.copy()
bgr, ratio, dwdh = letterbox(bgr, (W, H))
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
tensor = blob(rgb, return_seg=False)
dwdh = np.array(dwdh * 2, dtype=np.float32)
tensor = np.ascontiguousarray(tensor)
# inference
data = Engine(tensor)
print(f'image path {save_image}')
res = det_postprocess_unchanged_onnx(data)[0]
if len(res) == 0:
# if no bounding box
print(f'{image}: no object!')
continue
print(dwdh)
print(ratio)
# for (bbox, score, label) in zip(bboxes, scores, labels):
for i in range(len(res)):
bbox = res[i][:4]
label = res[i][-1]
score = res[i][4]
bbox -= dwdh
bbox /= ratio
bbox = bbox.round().astype(np.int32).tolist()
cls_id = int(label)
cls = CLASSES[cls_id]
color = COLORS[cls]
cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
cv2.putText(draw,
f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.75, [225, 255, 255],
thickness=2)
if args.show:
cv2.imshow('result', draw)
cv2.waitKey(0)
else:
cv2.imwrite(str(save_image), draw)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--engine', type=str, help='Engine file')
parser.add_argument('--imgs', type=str, help='Images file')
parser.add_argument('--show',
action='store_true',
help='Show the detection results')
parser.add_argument('--out-dir',
type=str,
default='./output',
help='Path to output file')
parser.add_argument('--method',
type=str,
default='cudart',
help='CUDART pipeline')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
main(args)
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/minwin/Yolov8_tensorrt.git
git@gitee.com:minwin/Yolov8_tensorrt.git
minwin
Yolov8_tensorrt
Yolov8_tensorrt
main

搜索帮助

0d507c66 1850385 C8b1a773 1850385