bbox.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. from __future__ import print_function
  2. import os
  3. import sys
  4. import cv2
  5. import random
  6. import datetime
  7. import time
  8. import math
  9. import argparse
  10. import numpy as np
  11. import torch
  12. try:
  13. from iou import IOU
  14. except BaseException:
  15. # IOU cython speedup 10x
  16. def IOU(ax1, ay1, ax2, ay2, bx1, by1, bx2, by2):
  17. sa = abs((ax2 - ax1) * (ay2 - ay1))
  18. sb = abs((bx2 - bx1) * (by2 - by1))
  19. x1, y1 = max(ax1, bx1), max(ay1, by1)
  20. x2, y2 = min(ax2, bx2), min(ay2, by2)
  21. w = x2 - x1
  22. h = y2 - y1
  23. if w < 0 or h < 0:
  24. return 0.0
  25. else:
  26. return 1.0 * w * h / (sa + sb - w * h)
  27. def bboxlog(x1, y1, x2, y2, axc, ayc, aww, ahh):
  28. xc, yc, ww, hh = (x2 + x1) / 2, (y2 + y1) / 2, x2 - x1, y2 - y1
  29. dx, dy = (xc - axc) / aww, (yc - ayc) / ahh
  30. dw, dh = math.log(ww / aww), math.log(hh / ahh)
  31. return dx, dy, dw, dh
  32. def bboxloginv(dx, dy, dw, dh, axc, ayc, aww, ahh):
  33. xc, yc = dx * aww + axc, dy * ahh + ayc
  34. ww, hh = math.exp(dw) * aww, math.exp(dh) * ahh
  35. x1, x2, y1, y2 = xc - ww / 2, xc + ww / 2, yc - hh / 2, yc + hh / 2
  36. return x1, y1, x2, y2
  37. def nms(dets, thresh):
  38. if 0 == len(dets):
  39. return []
  40. x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
  41. areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  42. order = scores.argsort()[::-1]
  43. keep = []
  44. while order.size > 0:
  45. i = order[0]
  46. keep.append(i)
  47. xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]])
  48. xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]])
  49. w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1)
  50. ovr = w * h / (areas[i] + areas[order[1:]] - w * h)
  51. inds = np.where(ovr <= thresh)[0]
  52. order = order[inds + 1]
  53. return keep
  54. def encode(matched, priors, variances):
  55. """Encode the variances from the priorbox layers into the ground truth boxes
  56. we have matched (based on jaccard overlap) with the prior boxes.
  57. Args:
  58. matched: (tensor) Coords of ground truth for each prior in point-form
  59. Shape: [num_priors, 4].
  60. priors: (tensor) Prior boxes in center-offset form
  61. Shape: [num_priors,4].
  62. variances: (list[float]) Variances of priorboxes
  63. Return:
  64. encoded boxes (tensor), Shape: [num_priors, 4]
  65. """
  66. # dist b/t match center and prior's center
  67. g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
  68. # encode variance
  69. g_cxcy /= (variances[0] * priors[:, 2:])
  70. # match wh / prior wh
  71. g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
  72. g_wh = torch.log(g_wh) / variances[1]
  73. # return target for smooth_l1_loss
  74. return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4]
  75. def decode(loc, priors, variances):
  76. """Decode locations from predictions using priors to undo
  77. the encoding we did for offset regression at train time.
  78. Args:
  79. loc (tensor): location predictions for loc layers,
  80. Shape: [num_priors,4]
  81. priors (tensor): Prior boxes in center-offset form.
  82. Shape: [num_priors,4].
  83. variances: (list[float]) Variances of priorboxes
  84. Return:
  85. decoded bounding box predictions
  86. """
  87. boxes = torch.cat((
  88. priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
  89. priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
  90. boxes[:, :2] -= boxes[:, 2:] / 2
  91. boxes[:, 2:] += boxes[:, :2]
  92. return boxes
  93. def batch_decode(loc, priors, variances):
  94. """Decode locations from predictions using priors to undo
  95. the encoding we did for offset regression at train time.
  96. Args:
  97. loc (tensor): location predictions for loc layers,
  98. Shape: [num_priors,4]
  99. priors (tensor): Prior boxes in center-offset form.
  100. Shape: [num_priors,4].
  101. variances: (list[float]) Variances of priorboxes
  102. Return:
  103. decoded bounding box predictions
  104. """
  105. boxes = torch.cat((
  106. priors[:, :, :2] + loc[:, :, :2] * variances[0] * priors[:, :, 2:],
  107. priors[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1])), 2)
  108. boxes[:, :, :2] -= boxes[:, :, 2:] / 2
  109. boxes[:, :, 2:] += boxes[:, :, :2]
  110. return boxes