RTDETRDetectionModel TorchScript, ONNX Predict and Val support (#8818)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
911d18e4f3
commit
af6c02c39b
7 changed files with 52 additions and 7 deletions
|
|
@ -94,6 +94,9 @@ class RTDETRValidator(DetectionValidator):
|
|||
|
||||
def postprocess(self, preds):
|
||||
"""Apply Non-maximum suppression to prediction outputs."""
|
||||
if not isinstance(preds, (list, tuple)): # list for PyTorch inference but list[0] Tensor for export inference
|
||||
preds = [preds, None]
|
||||
|
||||
bs, _, nd = preds[0].shape
|
||||
bboxes, scores = preds[0].split((4, nd - 4), dim=-1)
|
||||
bboxes *= self.args.imgsz
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue