bbox.py

from __future__ import print_function
import os,sys,cv2,random,datetime,time,math
import argparse
import numpy as np
import torch

try:
    from iou import IOU
except:
    # IOU cython speedup 10x
    def IOU(ax1,ay1,ax2,ay2,bx1,by1,bx2,by2):
        sa = abs((ax2-ax1)*(ay2-ay1))
        sb = abs((bx2-bx1)*(by2-by1))
        x1,y1 = max(ax1,bx1),max(ay1,by1)
        x2,y2 = min(ax2,bx2),min(ay2,by2)
        w  = x2 - x1
        h  = y2 - y1
        if w<0 or h<0: return 0.0
        else: return 1.0*w*h/(sa+sb-w*h)

def bboxlog(x1,y1,x2,y2,axc,ayc,aww,ahh):
    xc,yc,ww,hh = (x2+x1)/2,(y2+y1)/2,x2-x1,y2-y1
    dx,dy = (xc-axc)/aww,(yc-ayc)/ahh
    dw,dh = math.log(ww/aww),math.log(hh/ahh)
    return dx,dy,dw,dh

def bboxloginv(dx,dy,dw,dh,axc,ayc,aww,ahh):
    xc,yc = dx*aww+axc, dy*ahh+ayc
    ww,hh = math.exp(dw)*aww,math.exp(dh)*ahh
    x1,x2,y1,y2 = xc-ww/2,xc+ww/2,yc-hh/2,yc+hh/2
    return x1,y1,x2,y2

def nms(bboxlist:torch.Tensor, thresh:float) -> list:
    """Given an Nx5 tensor of bounding boxes, and a threshold,
    return a list of the indexes of bounding boxes to keep.
    """
    if len(bboxlist) == 0: 
        return []
    x1 = bboxlist[:,0]
    y1 = bboxlist[:,1]
    x2 = bboxlist[:,2]
    y2 = bboxlist[:,3]
    scores = bboxlist[:,4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    # Go through the boxes in order of decreasing score...
    scores = np.asarray(scores)
    order = scores.argsort()
    order = np.asarray(list(order[::-1]))
    keep = []
    while len(order) > 0:
        i = order[0]
        keep.append(i)  # Keep this one.

        # For all the remaining (lower score) bounding boxes, figure out something about the overlap I think
        xx1,yy1 = np.maximum(x1[i], x1[order[1:]]),np.maximum(y1[i], y1[order[1:]])
        xx2,yy2 = np.minimum(x2[i], x2[order[1:]]),np.minimum(y2[i], y2[order[1:]])
        w,h = np.maximum(0.0, xx2 - xx1 + 1),np.maximum(0.0, yy2 - yy1 + 1)
        ovr = w*h / (areas[i] + areas[order[1:]] - w*h)  # looks like the overlap
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]  # eliminate the ones that don't meet the threshhold

    return keep


def encode(matched, priors, variances):
    """Encode the variances from the priorbox layers into the ground truth boxes
    we have matched (based on jaccard overlap) with the prior boxes.
    Args:
        matched: (tensor) Coords of ground truth for each prior in point-form
            Shape: [num_priors, 4].
        priors: (tensor) Prior boxes in center-offset form
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        encoded boxes (tensor), Shape: [num_priors, 4]
    """

    # dist b/t match center and prior's center
    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
    # encode variance
    g_cxcy /= (variances[0] * priors[:, 2:])
    # match wh / prior wh
    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
    g_wh = torch.log(g_wh) / variances[1]
    # return target for smooth_l1_loss
    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]

def decode(loc, priors, variances):
    """Decode locations from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        loc (tensor): location predictions for loc layers,
            Shape: [num_priors,4]
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        decoded bounding box predictions
    """

    boxes = torch.cat((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes