diff --git a/README.md b/README.md index dcdabc1..6fb7d9a 100644 --- a/README.md +++ b/README.md @@ -2,12 +2,12 @@ YOLOv8 using TensorRT accelerate ! # Preprocessed ONNX model -[**YOLOv8-n**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8n_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772926549&Signature=xs8tFNWWZBaeHEwADlws4EYjov4%3D) +[**YOLOv8-n**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8n_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772936700&Signature=r6HgJTTcCSAxQxD9bKO9qBTtigQ%3D) -[**YOLOv8-s**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8s_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772926525&Signature=zfgKigShUuwcH%2BKhzrPs4Jmw%2FnE%3D) +[**YOLOv8-s**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8s_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1682936722&Signature=JjxQFx1YElcVdsCaMoj81KJ4a5s%3D) -[**YOLOv8-m**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8m_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1673926579&Signature=fe2LNWLzeWH3eS1UdqE9g49Anbs%3D) +[**YOLOv8-m**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8m_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1682936739&Signature=IRKBELdVFemD7diixxxgzMYqsWg%3D) -[**YOLOv8-l**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8l_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772926598&Signature=eD7Se4K6BZ9IgpaBx1tCNGkKEgk%3D) +[**YOLOv8-l**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8l_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1682936763&Signature=RGkJ4G2XJ4J%2BNiki5cJi3oBkDnA%3D) -[**YOLOv8-x**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8x_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772926618&Signature=xafM4XR7cikSbTlPyzrN%2FfxdiKM%3D) +[**YOLOv8-x**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8x_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1673936778&Signature=3o%2F7QKhiZg1dW3I6sDrY4ug6MQU%3D) diff --git a/data/bus.jpg b/data/bus.jpg new file mode 100644 index 0000000..b43e311 Binary files /dev/null and b/data/bus.jpg differ diff --git a/data/zidane.jpg b/data/zidane.jpg new file mode 100644 index 0000000..92d72ea Binary files /dev/null and b/data/zidane.jpg differ diff --git a/infer.py b/infer.py new file mode 100644 index 0000000..554900a --- /dev/null +++ b/infer.py @@ -0,0 +1,121 @@ +from models import TRTModule +from pathlib import Path +import cv2 +import argparse +import numpy as np +import torch +import random + +random.seed(0) + +SUFFIXS = ('.bmp', '.dng', '.jpeg', '.jpg', '.mpo', '.png', '.tif', '.tiff', '.webp', '.pfm') +CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', + 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', + 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', + 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', + 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', + 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', + 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', + 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', + 'hair drier', 'toothbrush') + +COLORS = {cls: [random.randint(0, 255) for _ in range(3)] for i, cls in enumerate(CLASSES)} + + +def letterbox(im, new_shape=(640, 640), color=(114, 114, 114)): + # Resize and pad image while meeting stride-multiple constraints + shape = im.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + # Scale ratio (new / old) + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + + # Compute padding + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding + + dw /= 2 # divide padding into 2 sides + dh /= 2 + + if shape[::-1] != new_unpad: # resize + im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border + return im, np.array([r, r, r, r], dtype=np.float32), np.array([dw, dh, dw, dh], dtype=np.float32) + + +def blob(im): + im = im.transpose(2, 0, 1) + im = im[np.newaxis, ...] + im = np.ascontiguousarray(im).astype(np.float32) / 255 + return im + + +def main(args): + device = torch.device(args.device) + Engine = TRTModule(args.engine, device) + + images_path = Path(args.imgs) + assert images_path.exists() + save_path = Path(args.out_dir) + + if images_path.is_dir(): + images = [i.absolute() for i in images_path.iterdir() if i.suffix in SUFFIXS] + else: + assert images_path.suffix in SUFFIXS + images = [images_path.absolute()] + + 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) + rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + tensor = blob(rgb) + ratio = torch.asarray(ratio, dtype=torch.float32, device=device) + dwdh = torch.asarray(dwdh, dtype=torch.float32, device=device) + tensor = torch.asarray(tensor, device=device) + num_dets, bboxes, scores, labels = Engine(tensor) + bboxes = bboxes[0, :num_dets.item()] + scores = scores[0, :num_dets.item()] + labels = labels[0, :num_dets.item()] + bboxes -= dwdh + bboxes /= ratio + for (bbox, score, label) in zip(bboxes, scores, labels): + bbox = bbox.round().int().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( + '--device', type=str, default='cuda:0', help='TensorRT infer device') + args = parser.parse_args() + return args + + +if __name__ == '__main__': + args = parse_args() + main(args) diff --git a/models/engine.py b/models/engine.py index 1f6c17d..e36cddd 100644 --- a/models/engine.py +++ b/models/engine.py @@ -119,6 +119,8 @@ class TRTModule(torch.nn.Module): if not dynamic: self.output_tensor = [torch.empty(info.shape, dtype=info.dtype, device=self.device) for info in out_info] self.is_dynamic = dynamic + self.inp_info = inp_info + self.out_infp = out_info def forward(self, *inputs):