Conversion to ONNX

#6
by MrIS-dotnet - opened

Thank you for the great work you have done to create your pre-trained model πŸ€—

I would like to ask if you have any experience or even can confirm that the RT-DETR model can be converted to ONNX and then on to TF Lite.

I tried to do it, I used torch.onnx.export, I got some results, but I am not able to perform the inference of this model, nor to evaluate it and find out the precision of the model.

Thank you in advance for your reply.

@MrIS-dotnet Thanks for the interest! I can not share the file but I can definitely say that you can convert into ONNX and TensoRT format.

@danelcsb Thank you for the quick response πŸ‘

Please, what tool did you use to convert to ONNX? I used torch.onnx.export, as I mentioned above with this list of arguments:

import torch
CHECKPOINT = "/content/rt-detr-finetuned/checkpoint-9300"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForObjectDetection.from_pretrained(CHECKPOINT, use_safetensors=True).to(DEVICE)
processor = AutoImageProcessor.from_pretrained(CHECKPOINT)
model.eval()
dummy_input = torch.randn(1, 3, 640, 640).to(DEVICE)
torch.onnx.export(
model,
dummy_input,
"rt-detr_model.onnx",
export_params=True,
opset_version=17,
input_names=["input"],
output_names=["output"]
)

I then checked the ONNX output model, no errors were found:

import onnx
onnx_model = onnx.load("rt-detr_model.onnx")
onnx.checker.check_model(onnx_model)

When I did an evaluation of the ONNX model, quality metrics like mAP resulted in 0 and I can't even make a prediction with the model.

import onnxruntime
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.patches as patches

Load the ONNX model

onnx_model_path = "/content/rt-detr_model.onnx"
ort_session = onnxruntime.InferenceSession(onnx_model_path, providers=["CPUExecutionProvider"])

Function to preprocess the image

def preprocess_image(image_path, input_shape=(640, 640)):
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, input_shape)
image = image.astype(np.float32) / 255.0
image = np.transpose(image, (2, 0, 1)) # Channels first
image = np.expand_dims(image, axis=0) # Add batch dimension
return image

Prepare the input

image_path = "/content/coco/valid/87acca98-scene00102.jpg"
input_tensor = preprocess_image(image_path)

Run inference

ort_inputs = {ort_session.get_inputs()[0].name: input_tensor}
ort_outs = ort_session.run(None, ort_inputs)

Inspect output shapes

for i, output in enumerate(ort_outs):
print(f"Output {i} shape: {output.shape}")

Example post-processing (adjust based on your model's output)

boxes = ort_outs[0][0] # Assuming first output is boxes
scores = ort_outs[1][0] # Assuming second output is scores
labels = ort_outs[2][0] # Assuming third output is labels

Draw bounding boxes on the image

image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
fig, ax = plt.subplots(1, figsize=(12, 9))
ax.imshow(image)

for box, score, label in zip(boxes, scores, labels):
if np.max(score) > 0.3: # Confidence threshold
x_min, y_min, x_max, y_max = box
width, height = x_max - x_min, y_max - y_min
rect = patches.Rectangle((x_min, y_min), width, height, linewidth=2, edgecolor='r', facecolor='none')
ax.add_patch(rect)
ax.text(x_min, y_min, f'{int(label)}: {score:.2f}', fontsize=12, color='white', bbox=dict(facecolor='red', alpha=0.5))

plt.axis('off')
plt.show()

Please, how did you proceed to export a working ONNX model from PyTorch RT-DETR?
Thank you.

@MrIS-dotnet There is nothing wrong in your code. Do you aware of postprocessing pipeline of ONNX does not include NMS?

As stated on the page here (https://docs.ultralytics.com/models/rtdetr/), RT-DETR is based on the idea of DETR (the NMS-free framework) so it really doesn't include NMS. Therefore even the ONNX model cannot contain NMS operations after export from RT-DETR.
Here https://github.com/NVIDIA/TensorRT/issues/795 I found some method how to add NMS operation to ONNX model in postprocessing, but it's for TensorRT and I have no experience with it, I don't know if it can be implemented in my scenario.
Could you please advise me how you did it?

@MrIS-dotnet Usually postprocessor is added after the training and export step. https://github.com/lyuwenyu/RT-DETR/blob/main/rtdetr_pytorch/tools/export_onnx.py
This is the concept of export_onnx from original code and the fundamental is not changed.

Sign up or log in to comment