search

lyuwenyu / RT-DETR

[CVPR 2024] Official RT-DETR (RTDETR paddle pytorch), Real-Time DEtection TRansformer, DETRs Beat YOLOs on Real-time Object Detection. 🔥 🔥 🔥
Apache License 2.0
2.61k stars 303 forks source link

:bug: Bugfix: TensorRT output types #395

Closed Hitbee-dev closed 3 months ago

Hitbee-dev commented 3 months ago

The output of labels, boxes, scores in rtdetrv2_tensorrt.py at line 228 should be changed as follows since they are of dictionary type.

Before:

labels, boxes, scores = m(blob)

draw([im_pil], labels, boxes, scores)

After:

output = m(blob)

draw([im_pil], output['labels'], output['boxes'], output['scores'])

Now, the RT-DETR model converted to TensorRT can produce recognition results correctly.