(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor] = None)
| 458 | |
| 459 | |
| 460 | def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor] = None): |
| 461 | cls_dist = cls_dist.sigmoid() * (1 if confidence is None else confidence) |
| 462 | |
| 463 | batch_idx, valid_grid, valid_cls = torch.where(cls_dist > nms_cfg.min_confidence) |
| 464 | valid_con = cls_dist[batch_idx, valid_grid, valid_cls] |
| 465 | valid_box = bbox[batch_idx, valid_grid] |
| 466 | |
| 467 | nms_idx = batched_nms(valid_box, valid_con, batch_idx + valid_cls * bbox.size(0), nms_cfg.min_iou) |
| 468 | predicts_nms = [] |
| 469 | for idx in range(cls_dist.size(0)): |
| 470 | instance_idx = nms_idx[idx == batch_idx[nms_idx]] |
| 471 | |
| 472 | predict_nms = torch.cat( |
| 473 | [valid_cls[instance_idx][:, None], valid_box[instance_idx], valid_con[instance_idx][:, None]], dim=-1 |
| 474 | ) |
| 475 | |
| 476 | predicts_nms.append(predict_nms[: nms_cfg.max_bbox]) |
| 477 | return predicts_nms |
| 478 | |
| 479 | |
| 480 | def calculate_map(predictions, ground_truths) -> Dict[str, Tensor]: |
no outgoing calls