🐛 Bugfix: TensorRT output types (#395)

Hitbee-dev · web-flow · commit 6b36128553b6 · 2024-07-29T16:31:23.000+08:00
diff --git a/rtdetrv2_pytorch/references/deploy/rtdetrv2_tensorrt.py b/rtdetrv2_pytorch/references/deploy/rtdetrv2_tensorrt.py
@@ -225,6 +225,6 @@ def draw(images, labels, boxes, scores, thrh = 0.6):
         'orig_target_sizes': orig_size.to(args.device),
     }
 
-    labels, boxes, scores = m(blob)
+    output = m(blob)
 
-    draw([im_pil], labels, boxes, scores)
+    draw([im_pil], output['labels'], output['boxes'], output['scores'])

Original file line number	Diff line number	Diff line change
`@@ -225,6 +225,6 @@ def draw(images, labels, boxes, scores, thrh = 0.6):`
`225`	`225`	`'orig_target_sizes': orig_size.to(args.device),`
`226`	`226`	`}`
`227`	`227`
`228`		`- labels, boxes, scores = m(blob)`
	`228`	`+ output = m(blob)`
`229`	`229`
`230`		`- draw([im_pil], labels, boxes, scores)`
	`230`	`+ draw([im_pil], output['labels'], output['boxes'], output['scores'])`