import torch


def yolo_filter_boxes(box_confidence: torch.Tensor, boxes: torch.Tensor, box_class_probs: torch.Tensor, threshold: float=.6):
    """filter boxes that has score smaller than threshold

    Parameters
    ----------
    box_confidence : torch.Tensor
        Object confidence
        shape [batch_size, num_anchors, 1, conv_height, conv_width]
    boxes : torch.Tensor
        boxes in format x_min, y_min, x_max, y_max
        shape [batch_size, num_anchors, 4, conv_height, conv_width]

    box_class_probs : torch.Tensor
        Class probalility
        shape [batch_size, num_anchors, num_classes, conv_height, conv_width]
    threshold : float, optional
        Score threshold

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        Filtered boxes, scores, classes
    """

    batch_size, num_anchors, _, conv_height, conv_width = box_confidence.shape

    box_scores = box_confidence * box_class_probs

    box_classes = torch.argmax(box_scores, dim=2, keepdim=True)

    box_class_scores, _ = torch.max(box_scores, dim=2, keepdim=True)

    prediction_mask = box_class_scores > threshold

    classes = box_classes[prediction_mask]
    scores = box_class_scores[prediction_mask]

    boxes = boxes.permute(0, 1, 3, 4, 2)
    prediction_mask = prediction_mask.permute(0, 1, 3, 4, 2)
    boxes = boxes[prediction_mask.expand_as(boxes)].view(-1, 4)

    return boxes, scores, classes