Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added the ability to get bounding boxes in output yolo when adding prepostprocessing #790

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Added the ability to get bounding boxes in output yolo when adding prepostprocessing #790

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
42ae06a
Merge pull request #1 from microsoft/main
VictoriaKlyueva Aug 13, 2024
4fe2185
Added option to get boxes in yolo output after prepostprocessing
VictoriaKlyueva Aug 13, 2024
1b86a28
Merge branch 'main' of https://github.com/VictoriaKlyueva/onnxruntime…
VictoriaKlyueva Aug 13, 2024
81db926
Merge branch 'main' into main
sayanshaw24 Aug 28, 2024
4b5dab9
Merge branch 'main' into main
VictoriaKlyueva Sep 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 18 additions & 12 deletions onnxruntime_extensions/tools/add_pre_post_processing_to_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,8 @@ def superresolution(model_file: Path, output_file: Path, output_format: str, onn


def yolo_detection(model_file: Path, output_file: Path, output_format: str = 'jpg',
onnx_opset: int = 16, num_classes: int = 80, input_shape: List[int] = None):
onnx_opset: int = 16, num_classes: int = 80, input_shape: List[int] = None,
output_as_image: bool = True):
"""
SSD-like model and Faster-RCNN-like model are including NMS inside already, You can find it from onnx model zoo.

Expand All @@ -185,6 +186,7 @@ def yolo_detection(model_file: Path, output_file: Path, output_format: str = 'jp
:param onnx_opset: The opset version of onnx model, default(16).
:param num_classes: The number of classes, default(80).
:param input_shape: The shape of input image (height,width), default will be asked from model input.
:param output_as_image: The flag that means that the model should have the image with boxes instead of the coordinates of the boxess
"""
model = onnx.load(str(model_file.resolve(strict=True)))
inputs = [create_named_value("image", onnx.TensorProto.UINT8, ["num_bytes"])]
Expand Down Expand Up @@ -284,19 +286,23 @@ def yolo_detection(model_file: Path, output_file: Path, output_format: str = 'jp
utils.IoMapEntry("Resize", producer_idx=0, consumer_idx=2),
utils.IoMapEntry("LetterBox", producer_idx=0, consumer_idx=3),
]),
# DrawBoundingBoxes on the original image
# Model imported from pytorch has CENTER_XYWH format
# two mode for how to color box,
# 1. colour_by_classes=True, (colour_by_classes), 2. colour_by_classes=False,(colour_by_confidence)
(DrawBoundingBoxes(mode='CENTER_XYWH', num_classes=num_classes, colour_by_classes=True),
[
utils.IoMapEntry("ConvertImageToBGR", producer_idx=0, consumer_idx=0),
utils.IoMapEntry("ScaleBoundingBoxes", producer_idx=0, consumer_idx=1),
]),
# Encode to jpg/png
ConvertBGRToImage(image_format=output_format),
]

if output_as_image:
post_processing_steps += [
# DrawBoundingBoxes on the original image
# Model imported from pytorch has CENTER_XYWH format
# two mode for how to color box,
# 1. colour_by_classes=True, (colour_by_classes), 2. colour_by_classes=False,(colour_by_confidence)
(DrawBoundingBoxes(mode='CENTER_XYWH', num_classes=num_classes, colour_by_classes=True),
[
utils.IoMapEntry("ConvertImageToBGR", producer_idx=0, consumer_idx=0),
utils.IoMapEntry("ScaleBoundingBoxes", producer_idx=0, consumer_idx=1),
]),
# Encode to jpg/png
ConvertBGRToImage(image_format=output_format),
]

pipeline.add_post_processing(post_processing_steps)

new_model = pipeline.run(model)
Expand Down