glenn-jocher pre-commit-ci[bot] commited on
Commit
601dbb8
1 Parent(s): bcc92e2

AutoAnchor improved initialization robustness (#6854)

Browse files

* Update AutoAnchor

* Update AutoAnchor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Files changed (1) hide show
  1. utils/autoanchor.py +10 -8
utils/autoanchor.py CHANGED
@@ -125,15 +125,17 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
125
  wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
126
  # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
127
 
128
- # Kmeans calculation
129
- LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
130
- s = wh.std(0) # sigmas for whitening
131
- k = kmeans(wh / s, n, iter=30)[0] * s # points
132
- if len(k) != n: # kmeans may return fewer points than requested if wh is insufficient or too similar
133
- LOGGER.warning(f'{PREFIX}WARNING: scipy.cluster.vq.kmeans returned only {len(k)} of {n} requested points')
 
 
 
134
  k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init
135
- wh = torch.tensor(wh, dtype=torch.float32) # filtered
136
- wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
137
  k = print_results(k, verbose=False)
138
 
139
  # Plot
 
125
  wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
126
  # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
127
 
128
+ # Kmeans init
129
+ try:
130
+ LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
131
+ assert n <= len(wh) # apply overdetermined constraint
132
+ s = wh.std(0) # sigmas for whitening
133
+ k = kmeans(wh / s, n, iter=30)[0] * s # points
134
+ assert n == len(k) # kmeans may return fewer points than requested if wh is insufficient or too similar
135
+ except Exception:
136
+ LOGGER.warning(f'{PREFIX}WARNING: switching strategies from kmeans to random init')
137
  k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init
138
+ wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0))
 
139
  k = print_results(k, verbose=False)
140
 
141
  # Plot