|
import math |
|
import sys |
|
from argparse import ArgumentParser |
|
from pathlib import Path |
|
|
|
import cv2 |
|
import onnxruntime |
|
from config import (CLASS_COLORS, CLASS_NAMES, ModelType, YOLOv5_ANCHORS, |
|
YOLOv7_ANCHORS) |
|
from cv2_nms import non_max_suppression |
|
from numpy_coder import Decoder |
|
from preprocess import Preprocess |
|
from tqdm import tqdm |
|
|
|
|
|
sys.path.append(str(Path(__file__).resolve().parents[0])) |
|
|
|
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', |
|
'.tiff', '.webp') |
|
|
|
|
|
def path_to_list(path: str): |
|
path = Path(path) |
|
if path.is_file() and path.suffix in IMG_EXTENSIONS: |
|
res_list = [str(path.absolute())] |
|
elif path.is_dir(): |
|
res_list = [ |
|
str(p.absolute()) for p in path.iterdir() |
|
if p.suffix in IMG_EXTENSIONS |
|
] |
|
else: |
|
raise RuntimeError |
|
return res_list |
|
|
|
|
|
def parse_args(): |
|
parser = ArgumentParser() |
|
parser.add_argument( |
|
'img', help='Image path, include image file, dir and URL.') |
|
parser.add_argument('onnx', type=str, help='Onnx file') |
|
parser.add_argument('--type', type=str, help='Model type') |
|
parser.add_argument( |
|
'--img-size', |
|
nargs='+', |
|
type=int, |
|
default=[640, 640], |
|
help='Image size of height and width') |
|
parser.add_argument( |
|
'--out-dir', default='./output', type=str, help='Path to output file') |
|
parser.add_argument( |
|
'--show', action='store_true', help='Show the detection results') |
|
parser.add_argument( |
|
'--score-thr', type=float, default=0.3, help='Bbox score threshold') |
|
parser.add_argument( |
|
'--iou-thr', type=float, default=0.7, help='Bbox iou threshold') |
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
def main(): |
|
args = parse_args() |
|
out_dir = Path(args.out_dir) |
|
model_type = ModelType(args.type.lower()) |
|
|
|
if not args.show: |
|
out_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
files = path_to_list(args.img) |
|
session = onnxruntime.InferenceSession( |
|
args.onnx, providers=['CPUExecutionProvider']) |
|
preprocessor = Preprocess(model_type) |
|
decoder = Decoder(model_type, model_only=True) |
|
if model_type == ModelType.YOLOV5: |
|
anchors = YOLOv5_ANCHORS |
|
elif model_type == ModelType.YOLOV7: |
|
anchors = YOLOv7_ANCHORS |
|
else: |
|
anchors = None |
|
|
|
for file in tqdm(files): |
|
image = cv2.imread(file) |
|
image_h, image_w = image.shape[:2] |
|
img, (ratio_w, ratio_h) = preprocessor(image, args.img_size) |
|
features = session.run(None, {'images': img}) |
|
decoder_outputs = decoder( |
|
features, |
|
args.score_thr, |
|
num_labels=len(CLASS_NAMES), |
|
anchors=anchors) |
|
nmsd_boxes, nmsd_scores, nmsd_labels = non_max_suppression( |
|
*decoder_outputs, args.score_thr, args.iou_thr) |
|
for box, score, label in zip(nmsd_boxes, nmsd_scores, nmsd_labels): |
|
x0, y0, x1, y1 = box |
|
x0 = math.floor(min(max(x0 / ratio_w, 1), image_w - 1)) |
|
y0 = math.floor(min(max(y0 / ratio_h, 1), image_h - 1)) |
|
x1 = math.ceil(min(max(x1 / ratio_w, 1), image_w - 1)) |
|
y1 = math.ceil(min(max(y1 / ratio_h, 1), image_h - 1)) |
|
cv2.rectangle(image, (x0, y0), (x1, y1), CLASS_COLORS[label], 2) |
|
cv2.putText(image, f'{CLASS_NAMES[label]}: {score:.2f}', |
|
(x0, y0 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, |
|
(0, 255, 255), 2) |
|
if args.show: |
|
cv2.imshow('result', image) |
|
cv2.waitKey(0) |
|
else: |
|
cv2.imwrite(f'{out_dir / Path(file).name}', image) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
