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

Faster IoU prediction matching by removing torch.cat (#4708)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Andy 2023-09-03 18:19:30 +02:00 committed by GitHub
parent 2ba80e355a
commit 02b857e14c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 25 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
ultralytics/engine/validator.py
View File

@ -24,6 +24,7 @@ from pathlib import Path
import numpy as np import numpy as np
import torch import torch
from scipy.optimize import linear_sum_assignment
from ultralytics.cfg import get_cfg, get_save_dir from ultralytics.cfg import get_cfg, get_save_dir
from ultralytics.data.utils import check_cls_dataset, check_det_dataset from ultralytics.data.utils import check_cls_dataset, check_det_dataset
@ -204,7 +205,7 @@ class BaseValidator:
LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}") LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}")
return stats return stats
def match_predictions(self, pred_classes, true_classes, iou): def match_predictions(self, pred_classes, true_classes, iou, use_scipy=False):
""" """
Matches predictions to ground truth objects (pred_classes, true_classes) using IoU. Matches predictions to ground truth objects (pred_classes, true_classes) using IoU.
@ -212,23 +213,35 @@ class BaseValidator:
pred_classes (torch.Tensor): Predicted class indices of shape(N,). pred_classes (torch.Tensor): Predicted class indices of shape(N,).
true_classes (torch.Tensor): Target class indices of shape(M,). true_classes (torch.Tensor): Target class indices of shape(M,).
iou (torch.Tensor): An NxM tensor containing the pairwise IoU values for predictions and ground of truth iou (torch.Tensor): An NxM tensor containing the pairwise IoU values for predictions and ground of truth
use_scipy (bool): Whether to use scipy for matching (more precise).
Returns: Returns:
(torch.Tensor): Correct tensor of shape(N,10) for 10 IoU thresholds. (torch.Tensor): Correct tensor of shape(N,10) for 10 IoU thresholds.
""" """
# Dx10 matrix, where D - detections, 10 - IoU thresholds
correct = np.zeros((pred_classes.shape[0], self.iouv.shape[0])).astype(bool) correct = np.zeros((pred_classes.shape[0], self.iouv.shape[0])).astype(bool)
# LxD matrix where L - labels (rows), D - detections (columns)
correct_class = true_classes[:, None] == pred_classes correct_class = true_classes[:, None] == pred_classes
for i, iouv in enumerate(self.iouv): iou = iou * correct_class # zero out the wrong classes
x = torch.nonzero(iou.ge(iouv) & correct_class) # IoU > threshold and classes match iou = iou.cpu().numpy()
if x.shape[0]: for i, threshold in enumerate(self.iouv.cpu().tolist()):
# Concatenate [label, detect, iou] if use_scipy:
matches = torch.cat((x, iou[x[:, 0], x[:, 1]].unsqueeze(1)), 1).cpu().numpy() cost_matrix = iou * (iou >= threshold)
if x.shape[0] > 1: if cost_matrix.any():
matches = matches[matches[:, 2].argsort()[::-1]] labels_idx, detections_idx = linear_sum_assignment(cost_matrix, maximize=True)
matches = matches[np.unique(matches[:, 1], return_index=True)[1]] valid = cost_matrix[labels_idx, detections_idx] > 0
# matches = matches[matches[:, 2].argsort()[::-1]] if valid.any():
matches = matches[np.unique(matches[:, 0], return_index=True)[1]] correct[detections_idx[valid], i] = True
correct[matches[:, 1].astype(int), i] = True else:
matches = np.nonzero(iou >= threshold) # IoU > threshold and classes match
matches = np.array(matches).T
if matches.shape[0]:
if matches.shape[0] > 1:
matches = matches[iou[matches[:, 0], matches[:, 1]].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
# matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
return torch.tensor(correct, dtype=torch.bool, device=pred_classes.device) return torch.tensor(correct, dtype=torch.bool, device=pred_classes.device)
def add_callback(self, event: str, callback): def add_callback(self, event: str, callback):