|
- # --------------------------------------------------------
- # Modified from https://github.com/biubug6/Pytorch_Retinaface
- # --------------------------------------------------------
-
- from itertools import product as product
- from math import ceil
-
- import numpy as np
- import torch
-
-
- class PriorBox(object):
-
- def __init__(self, cfg, image_size=None, phase='train'):
- super(PriorBox, self).__init__()
- self.min_sizes = cfg['min_sizes']
- self.steps = cfg['steps']
- self.clip = cfg['clip']
- self.image_size = image_size
- self.feature_maps = [[
- ceil(self.image_size[0] / step),
- ceil(self.image_size[1] / step)
- ] for step in self.steps]
- self.name = 's'
-
- def forward(self):
- anchors = []
- for k, f in enumerate(self.feature_maps):
- min_sizes = self.min_sizes[k]
- for i, j in product(range(f[0]), range(f[1])):
- for min_size in min_sizes:
- s_kx = min_size / self.image_size[1]
- s_ky = min_size / self.image_size[0]
- dense_cx = [
- x * self.steps[k] / self.image_size[1]
- for x in [j + 0.5]
- ]
- dense_cy = [
- y * self.steps[k] / self.image_size[0]
- for y in [i + 0.5]
- ]
- for cy, cx in product(dense_cy, dense_cx):
- anchors += [cx, cy, s_kx, s_ky]
-
- # back to torch land
- output = torch.Tensor(anchors).view(-1, 4)
- if self.clip:
- output.clamp_(max=1, min=0)
- return output
-
-
- def py_cpu_nms(dets, thresh):
- """Pure Python NMS baseline."""
- x1 = dets[:, 0]
- y1 = dets[:, 1]
- x2 = dets[:, 2]
- y2 = dets[:, 3]
- scores = dets[:, 4]
-
- areas = (x2 - x1 + 1) * (y2 - y1 + 1)
- order = scores.argsort()[::-1]
-
- keep = []
- while order.size > 0:
- i = order[0]
- keep.append(i)
- xx1 = np.maximum(x1[i], x1[order[1:]])
- yy1 = np.maximum(y1[i], y1[order[1:]])
- xx2 = np.minimum(x2[i], x2[order[1:]])
- yy2 = np.minimum(y2[i], y2[order[1:]])
-
- w = np.maximum(0.0, xx2 - xx1 + 1)
- h = np.maximum(0.0, yy2 - yy1 + 1)
- inter = w * h
- ovr = inter / (areas[i] + areas[order[1:]] - inter)
-
- inds = np.where(ovr <= thresh)[0]
- order = order[inds + 1]
-
- return keep
-
-
- # Adapted from https://github.com/Hakuyume/chainer-ssd
- 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
-
-
- def decode_landm(pre, priors, variances):
- """Decode landm from predictions using priors to undo
- the encoding we did for offset regression at train time.
- Args:
- pre (tensor): landm predictions for loc layers,
- Shape: [num_priors,10]
- priors (tensor): Prior boxes in center-offset form.
- Shape: [num_priors,4].
- variances: (list[float]) Variances of priorboxes
- Return:
- decoded landm predictions
- """
- a = priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:]
- b = priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:]
- c = priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:]
- d = priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:]
- e = priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:]
- landms = torch.cat((a, b, c, d, e), dim=1)
- return landms
|
