box_utils.py

from typing import Tuple, Any 
import numpy as np
from savant.base.converter import BaseComplexModelOutputConverter
from savant.base.model import ComplexModel
from numba.typed import List

from savant.selector.detector import nms_cpu


def area_of(left_top, right_bottom):
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N, 2): left top corner.
        right_bottom (N, 2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    hw = np.clip(right_bottom - left_top, a_min=0.0, a_max=None)
    return hw[..., 0] * hw[..., 1]


def iou_of(boxes0, boxes1, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        boxes0 (N, 4): ground truth boxes.
        boxes1 (N or 1, 4): predicted boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
    overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])

    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)


def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
    """

    Args:
        box_scores (N, 5): boxes in corner-form and probabilities.
        iou_threshold: intersection over union threshold.
        top_k: keep top_k results. If k <= 0, keep all the results.
        candidate_size: only consider the candidates with the highest scores.
    Returns:
         picked: a list of indexes of the kept boxes
    """
    scores = box_scores[:, -1]
    boxes = box_scores[:, :-1]
    picked = []
    indexes = np.argsort(scores)[::-1]
    indexes = indexes[:candidate_size]
    while len(indexes) > 0:
        current = indexes[0]
        picked.append(current)
        if 0 < top_k == len(picked) or len(indexes) == 1:
            break
        current_box = boxes[current, :]
        indexes = indexes[1:]
        rest_boxes = boxes[indexes, :]
        iou = iou_of(
            rest_boxes,
            np.expand_dims(current_box, axis=0),
        )
        indexes = indexes[iou <= iou_threshold]

    return box_scores[picked, :]


class FaceConverter(BaseComplexModelOutputConverter):
  def __init__(
      self,
      confidence_threshold: float = 0.6,
      nms_iou_threshold: float = 0.5,
      **kwargs
  ):  
      super().__init__(**kwargs)
      self.confidence_threshold = confidence_threshold
      self.nms_iou_threshold = nms_iou_threshold

  def __call__(
      self,
      *output_layers: np.ndarray,
      model: ComplexModel,
      roi: Tuple[float, float, float, float],
  ) -> Tuple[np.ndarray, List[List[Tuple[str, Any, float]]]]:
      confidences, boxes = output_layers
      boxes = np.array(boxes[0])
      confidences = np.array(confidences[0])
      picked_box_probs = []
      picked_labels = []
      for class_index in range(1, confidences.shape[1]):
          probs = confidences[:, class_index]
          mask = probs > prob_threshold
          probs = probs[mask]
          if probs.shape[0] == 0:
              continue
          subset_boxes = boxes[mask, :]
              box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)
              box_probs = box_utils.hard_nms(box_probs,
              iou_threshold=iou_threshold,
              top_k=top_k,
          )   
          picked_box_probs.append(box_probs)
          picked_labels.extend([class_index] * box_probs.shape[0])
      if not picked_box_probs:
          return np.array([]), np.array([]), np.array([])

      picked_box_probs = np.concatenate(picked_box_probs)
      picked_box_probs[:, 0] *= width
      picked_box_probs[:, 1] *= height
      picked_box_probs[:, 2] *= width
      return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]