def process_centernet_output(
    predicted_heatmap,  # logits
    predicted_regression,
    gt_boxes,
    gt_labels,
    confidence_threshold=0.5,
    iou_threshold=0.5,
):
    """Generate bbox and classes from CenterNet model outputs.

    Args:
        predicted_heatmap (torch.Tensor): predicted center heatmap logits,
            expected shapes [batch, height, width, num classes].
        predicted_regression (torch.Tensor): predicted HW regression,
            expected shapes [batch, height, width, 2].
        gt_boxes (List[torch.Tensor]): list with sample bounding boxes.
        gt_labels (List[torch.Tensor]): list with sample bounding box labels.
        confidence_threshold (float): confidence threshold,
            proposals with lover values than threshold will be ignored.
            Default is ``0.5``.
        iou_threshold (float): IoU threshold to use in NMS.
            Default is ``0.5``.

    Yields:
        predicted sample (np.ndarray) and ground truth sample (np.ndarray)
    """
    batch_size = predicted_heatmap.size(0)

    hm = predicted_heatmap.sigmoid()
    pooled = F.max_pool2d(hm, kernel_size=(3, 3), stride=1, padding=1)
    hm *= torch.logical_and(
        hm >= confidence_threshold, pooled >= confidence_threshold
    ).float()

    hm_numpy = hm.detach().cpu().numpy()
    reg_numpy = predicted_regression.detach().cpu().numpy()

    for i in range(batch_size):
        sample_boxes = []
        sample_classes = []
        sample_scores = []
        for cls_idx in range(hm_numpy.shape[1]):
            # build predictions
            cls_boxes, cls_scores = pred2box(
                hm_numpy[i, cls_idx], reg_numpy[i], threshold=0, scale=4, input_size=512
            )

            # skip empty label predictions
            if cls_scores.shape[0] == 0:
                continue

            cls_boxes = cls_boxes / 512.0

            cls_boxes, cls_classes, cls_scores = nms_filter(
                cls_boxes,
                np.full(len(cls_scores), cls_idx),
                cls_scores,
                iou_threshold=iou_threshold,
            )
            sample_boxes.append(cls_boxes)
            sample_classes.append(cls_classes)
            sample_scores.append(cls_scores)
        # skip empty predictions
        if len(sample_boxes) == 0:
            continue

        sample_boxes = np.concatenate(sample_boxes, 0)
        sample_classes = np.concatenate(sample_classes, 0)
        sample_scores = np.concatenate(sample_scores, 0)

        pred_sample = np.concatenate(
            [sample_boxes, sample_classes[:, None], sample_scores[:, None]], -1
        )
        pred_sample = pred_sample.astype(np.float32)

        sample_gt_bboxes = gt_boxes[i].detach().cpu()
        sample_gt_classes = gt_labels[i].detach().cpu()
        gt_sample = np.zeros((sample_gt_classes.shape[0], 7), dtype=np.float32)
        gt_sample[:, :4] = sample_gt_bboxes
        gt_sample[:, 4] = sample_gt_classes

        yield pred_sample, gt_sample
Generate bbox and classes from CenterNet model outputs.

    Args:
        predicted_heatmap (torch.Tensor): predicted center heatmap logits,
            expected shapes [batch, height, width, num classes].
        predicted_regression (torch.Tensor): predicted HW regression,
            expected shapes [batch, height, width, 2].
        gt_boxes (List[torch.Tensor]): list with sample bounding boxes.
        gt_labels (List[torch.Tensor]): list with sample bounding box labels.
        confidence_threshold (float): confidence threshold,
            proposals with lover values than threshold will be ignored.
            Default is ``0.5``.
        iou_threshold (float): IoU threshold to use in NMS.
            Default is ``0.5``.

    Yields:
        predicted sample (np.ndarray) and ground truth sample (np.ndarray)
    
examples/detection/callbacks.py