Upload onnx_inference.py

onnx_inference.py

@@ -0,0 +1,160 @@
+import argparse
+import os
+
+import cv2
+import numpy as np
+from loguru import logger
+
+import onnxruntime
+
+from yolox.data.data_augment import preproc as preprocess
+from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis
+from yolox.utils.visualize import plot_tracking
+from yolox.tracker.byte_tracker import BYTETracker
+from yolox.tracking_utils.timer import Timer
+
+
+def make_parser():
+    parser = argparse.ArgumentParser("onnxruntime inference sample")
+    parser.add_argument(
+        "-m",
+        "--model",
+        type=str,
+        default="bytetrack_s.onnx",
+        help="Input your onnx model.",
+    )
+    parser.add_argument(
+        "-i",
+        "--video_path",
+        type=str,
+        default='../../videos/palace.mp4',
+        help="Path to your input image.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        type=str,
+        default='demo_output',
+        help="Path to your output directory.",
+    )
+    parser.add_argument(
+        "-s",
+        "--score_thr",
+        type=float,
+        default=0.1,
+        help="Score threshould to filter the result.",
+    )
+    parser.add_argument(
+        "-n",
+        "--nms_thr",
+        type=float,
+        default=0.7,
+        help="NMS threshould.",
+    )
+    parser.add_argument(
+        "--input_shape",
+        type=str,
+        default="608,1088",
+        help="Specify an input shape for inference.",
+    )
+    parser.add_argument(
+        "--with_p6",
+        action="store_true",
+        help="Whether your model uses p6 in FPN/PAN.",
+    )
+    # tracking args
+    parser.add_argument("--track_thresh", type=float, default=0.5, help="tracking confidence threshold")
+    parser.add_argument("--track_buffer", type=int, default=30, help="the frames for keep lost tracks")
+    parser.add_argument("--match_thresh", type=int, default=0.8, help="matching threshold for tracking")
+    parser.add_argument('--min-box-area', type=float, default=10, help='filter out tiny boxes')
+    parser.add_argument("--mot20", dest="mot20", default=False, action="store_true", help="test mot20.")
+    return parser
+
+
+class Predictor(object):
+    def __init__(self, args):
+        self.rgb_means = (0.485, 0.456, 0.406)
+        self.std = (0.229, 0.224, 0.225)
+        self.args = args
+        self.session = onnxruntime.InferenceSession(args.model)
+        self.input_shape = tuple(map(int, args.input_shape.split(',')))
+
+    def inference(self, ori_img, timer):
+        img_info = {"id": 0}
+        height, width = ori_img.shape[:2]
+        img_info["height"] = height
+        img_info["width"] = width
+        img_info["raw_img"] = ori_img
+
+        img, ratio = preprocess(ori_img, self.input_shape, self.rgb_means, self.std)
+        img_info["ratio"] = ratio
+        ort_inputs = {self.session.get_inputs()[0].name: img[None, :, :, :]}
+        timer.tic()
+        output = self.session.run(None, ort_inputs)
+        predictions = demo_postprocess(output[0], self.input_shape, p6=self.args.with_p6)[0]
+
+        boxes = predictions[:, :4]
+        scores = predictions[:, 4:5] * predictions[:, 5:]
+
+        boxes_xyxy = np.ones_like(boxes)
+        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
+        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
+        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
+        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
+        boxes_xyxy /= ratio
+        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=self.args.nms_thr, score_thr=self.args.score_thr)
+        return dets[:, :-1], img_info
+
+
+def imageflow_demo(predictor, args):
+    cap = cv2.VideoCapture(args.video_path)
+    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float
+    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    save_folder = args.output_dir
+    os.makedirs(save_folder, exist_ok=True)
+    save_path = os.path.join(save_folder, args.video_path.split("/")[-1])
+    logger.info(f"video save_path is {save_path}")
+    vid_writer = cv2.VideoWriter(
+        save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (int(width), int(height))
+    )
+    tracker = BYTETracker(args, frame_rate=30)
+    timer = Timer()
+    frame_id = 0
+    results = []
+    while True:
+        if frame_id % 20 == 0:
+            logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1. / max(1e-5, timer.average_time)))
+        ret_val, frame = cap.read()
+        if ret_val:
+            outputs, img_info = predictor.inference(frame, timer)
+            online_targets = tracker.update(outputs, [img_info['height'], img_info['width']], [img_info['height'], img_info['width']])
+            online_tlwhs = []
+            online_ids = []
+            online_scores = []
+            for t in online_targets:
+                tlwh = t.tlwh
+                tid = t.track_id
+                vertical = tlwh[2] / tlwh[3] > 1.6
+                if tlwh[2] * tlwh[3] > args.min_box_area and not vertical:
+                    online_tlwhs.append(tlwh)
+                    online_ids.append(tid)
+                    online_scores.append(t.score)
+            timer.toc()
+            results.append((frame_id + 1, online_tlwhs, online_ids, online_scores))
+            online_im = plot_tracking(img_info['raw_img'], online_tlwhs, online_ids, frame_id=frame_id + 1,
+                                      fps=1. / timer.average_time)
+            vid_writer.write(online_im)
+            ch = cv2.waitKey(1)
+            if ch == 27 or ch == ord("q") or ch == ord("Q"):
+                break
+        else:
+            break
+        frame_id += 1
+
+
+if __name__ == '__main__':
+    args = make_parser().parse_args()
+
+    predictor = Predictor(args)
+    imageflow_demo(predictor, args)