HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.54 New Streamlit inference Solution (#18316)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-12-24 16:26:56 +05:00 • committed by GitHub
parent 5b76bed7d0
commit 51026a9a4a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 251 additions and 188 deletions
Download patch file Download diff file Expand all files Collapse all files

23
examples/RTDETR-ONNXRuntime-Python/main.py
View file

@ -8,13 +8,14 @@ import torch
from ultralytics.utils import ASSETS, yaml_load
from ultralytics.utils.checks import check_requirements, check_yaml
class RTDETR:
"""RTDETR object detection model class for handling inference and visualization."""
def __init__(self, model_path, img_path, conf_thres=0.5, iou_thres=0.5):
"""
Initializes the RTDETR object with the specified parameters.
Args:
model_path: Path to the ONNX model file.
img_path: Path to the input image.
@ -71,11 +72,17 @@ class RTDETR:
# Draw a filled rectangle as the background for the label text
cv2.rectangle(
self.img, (int(label_x), int(label_y - label_height)), (int(label_x + label_width), int(label_y + label_height)), color, cv2.FILLED
self.img,
(int(label_x), int(label_y - label_height)),
(int(label_x + label_width), int(label_y + label_height)),
color,
cv2.FILLED,
)
# Draw the label text on the image
cv2.putText(self.img, label, (int(label_x), int(label_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
cv2.putText(
self.img, label, (int(label_x), int(label_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA
)
def preprocess(self):
"""
@ -110,8 +117,7 @@ class RTDETR:
def bbox_cxcywh_to_xyxy(self, boxes):
"""
Converts bounding boxes from (center x, center y, width, height) format
to (x_min, y_min, x_max, y_max) format.
Converts bounding boxes from (center x, center y, width, height) format to (x_min, y_min, x_max, y_max) format.
Args:
boxes (numpy.ndarray): An array of shape (N, 4) where each row represents
@ -176,7 +182,7 @@ class RTDETR:
def main(self):
"""
Executes the detection on the input image using the ONNX model.
Returns:
np.array: Output image with annotations.
"""
@ -189,6 +195,7 @@ class RTDETR:
# Process and return the model output
return self.postprocess(model_output)
if __name__ == "__main__":
# Set up argument parser for command-line arguments
parser = argparse.ArgumentParser()
@ -210,4 +217,4 @@ if __name__ == "__main__":
# Display the annotated output image
cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
cv2.imshow("Output", output_image)
cv2.waitKey(0)
cv2.waitKey(0)