xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add Non-Maximum Suppression (NMS) inplace flag (#8368)

Browse Source
This commit is contained in:
AlainSchoebi 2024-02-23 14:28:25 +01:00 committed by GitHub
parent 21088e3986
commit d6edf1c1d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ultralytics/utils/ops.py
View File

@ -173,6 +173,7 @@ def non_max_suppression(
max_time_img=0.05,
max_nms=30000,
max_wh=7680,
in_place=True,
rotated=False,
):
"""
@ -197,7 +198,8 @@ def non_max_suppression(
nc (int, optional): The number of classes output by the model. Any indices after this will be considered masks.
max_time_img (float): The maximum time (seconds) for processing one image.
max_nms (int): The maximum number of boxes into torchvision.ops.nms().
max_wh (int): The maximum box width and height in pixels
max_wh (int): The maximum box width and height in pixels.
in_place (bool): If True, the input prediction tensor will be modified in place.
Returns:
(List[torch.Tensor]): A list of length batch_size, where each element is a tensor of
@ -224,7 +226,10 @@ def non_max_suppression(
prediction = prediction.transpose(-1, -2) # shape(1,84,6300) to shape(1,6300,84)
if not rotated:
prediction[..., :4] = xywh2xyxy(prediction[..., :4]) # xywh to xyxy
if in_place:
prediction[..., :4] = xywh2xyxy(prediction[..., :4]) # xywh to xyxy
else:
prediction = torch.cat((xywh2xyxy(prediction[..., :4]), prediction[..., 4:]), dim=-1) # xywh to xyxy
t = time.time()
output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs