nn_matching.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177
  1. # vim: expandtab:ts=4:sw=4
  2. import numpy as np
  3. def _pdist(a, b):
  4. """Compute pair-wise squared distance between points in `a` and `b`.
  5. Parameters
  6. ----------
  7. a : array_like
  8. An NxM matrix of N samples of dimensionality M.
  9. b : array_like
  10. An LxM matrix of L samples of dimensionality M.
  11. Returns
  12. -------
  13. ndarray
  14. Returns a matrix of size len(a), len(b) such that eleement (i, j)
  15. contains the squared distance between `a[i]` and `b[j]`.
  16. """
  17. a, b = np.asarray(a), np.asarray(b)
  18. if len(a) == 0 or len(b) == 0:
  19. return np.zeros((len(a), len(b)))
  20. a2, b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1)
  21. r2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :]
  22. r2 = np.clip(r2, 0., float(np.inf))
  23. return r2
  24. def _cosine_distance(a, b, data_is_normalized=False):
  25. """Compute pair-wise cosine distance between points in `a` and `b`.
  26. Parameters
  27. ----------
  28. a : array_like
  29. An NxM matrix of N samples of dimensionality M.
  30. b : array_like
  31. An LxM matrix of L samples of dimensionality M.
  32. data_is_normalized : Optional[bool]
  33. If True, assumes rows in a and b are unit length vectors.
  34. Otherwise, a and b are explicitly normalized to lenght 1.
  35. Returns
  36. -------
  37. ndarray
  38. Returns a matrix of size len(a), len(b) such that eleement (i, j)
  39. contains the squared distance between `a[i]` and `b[j]`.
  40. """
  41. if not data_is_normalized:
  42. a = np.asarray(a) / np.linalg.norm(a, axis=1, keepdims=True)
  43. b = np.asarray(b) / np.linalg.norm(b, axis=1, keepdims=True)
  44. return 1. - np.dot(a, b.T)
  45. def _nn_euclidean_distance(x, y):
  46. """ Helper function for nearest neighbor distance metric (Euclidean).
  47. Parameters
  48. ----------
  49. x : ndarray
  50. A matrix of N row-vectors (sample points).
  51. y : ndarray
  52. A matrix of M row-vectors (query points).
  53. Returns
  54. -------
  55. ndarray
  56. A vector of length M that contains for each entry in `y` the
  57. smallest Euclidean distance to a sample in `x`.
  58. """
  59. distances = _pdist(x, y)
  60. return np.maximum(0.0, distances.min(axis=0))
  61. def _nn_cosine_distance(x, y):
  62. """ Helper function for nearest neighbor distance metric (cosine).
  63. Parameters
  64. ----------
  65. x : ndarray
  66. A matrix of N row-vectors (sample points).
  67. y : ndarray
  68. A matrix of M row-vectors (query points).
  69. Returns
  70. -------
  71. ndarray
  72. A vector of length M that contains for each entry in `y` the
  73. smallest cosine distance to a sample in `x`.
  74. """
  75. distances = _cosine_distance(x, y)
  76. return distances.min(axis=0)
  77. class NearestNeighborDistanceMetric(object):
  78. """
  79. A nearest neighbor distance metric that, for each target, returns
  80. the closest distance to any sample that has been observed so far.
  81. Parameters
  82. ----------
  83. metric : str
  84. Either "euclidean" or "cosine".
  85. matching_threshold: float
  86. The matching threshold. Samples with larger distance are considered an
  87. invalid match.
  88. budget : Optional[int]
  89. If not None, fix samples per class to at most this number. Removes
  90. the oldest samples when the budget is reached.
  91. Attributes
  92. ----------
  93. samples : Dict[int -> List[ndarray]]
  94. A dictionary that maps from target identities to the list of samples
  95. that have been observed so far.
  96. """
  97. def __init__(self, metric, matching_threshold, budget=None):
  98. if metric == "euclidean":
  99. self._metric = _nn_euclidean_distance
  100. elif metric == "cosine":
  101. self._metric = _nn_cosine_distance
  102. else:
  103. raise ValueError(
  104. "Invalid metric; must be either 'euclidean' or 'cosine'")
  105. self.matching_threshold = matching_threshold
  106. self.budget = budget
  107. self.samples = {}
  108. def partial_fit(self, features, targets, active_targets):
  109. """Update the distance metric with new data.
  110. Parameters
  111. ----------
  112. features : ndarray
  113. An NxM matrix of N features of dimensionality M.
  114. targets : ndarray
  115. An integer array of associated target identities.
  116. active_targets : List[int]
  117. A list of targets that are currently present in the scene.
  118. """
  119. for feature, target in zip(features, targets):
  120. self.samples.setdefault(target, []).append(feature)
  121. if self.budget is not None:
  122. self.samples[target] = self.samples[target][-self.budget:]
  123. self.samples = {k: self.samples[k] for k in active_targets}
  124. def distance(self, features, targets):
  125. """Compute distance between features and targets.
  126. Parameters
  127. ----------
  128. features : ndarray
  129. An NxM matrix of N features of dimensionality M.
  130. targets : List[int]
  131. A list of targets to match the given `features` against.
  132. Returns
  133. -------
  134. ndarray
  135. Returns a cost matrix of shape len(targets), len(features), where
  136. element (i, j) contains the closest squared distance between
  137. `targets[i]` and `features[j]`.
  138. """
  139. cost_matrix = np.zeros((len(targets), len(features)))
  140. for i, target in enumerate(targets):
  141. cost_matrix[i, :] = self._metric(self.samples[target], features)
  142. return cost_matrix