glenn-jocher commited on
Commit
2e5c67e
β€’
1 Parent(s): 7b80545

Robust `scipy.cluster.vq.kmeans` too few points (#6668)

Browse files

* Handle `scipy.cluster.vq.kmeans` too few points

Resolves #6664

* Update autoanchor.py

* Cleanup

Files changed (1) hide show
  1. utils/autoanchor.py +7 -6
utils/autoanchor.py CHANGED
@@ -1,6 +1,6 @@
1
  # YOLOv5 πŸš€ by Ultralytics, GPL-3.0 license
2
  """
3
- Auto-anchor utils
4
  """
5
 
6
  import random
@@ -81,6 +81,7 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
81
  """
82
  from scipy.cluster.vq import kmeans
83
 
 
84
  thr = 1 / thr
85
 
86
  def metric(k, wh): # compute metrics
@@ -121,14 +122,15 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
121
  if i:
122
  LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
123
  wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
124
- # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
125
 
126
  # Kmeans calculation
127
  LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
128
  s = wh.std(0) # sigmas for whitening
129
- k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
130
- assert len(k) == n, f'{PREFIX}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}'
131
- k *= s
 
132
  wh = torch.tensor(wh, dtype=torch.float32) # filtered
133
  wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
134
  k = print_results(k, verbose=False)
@@ -146,7 +148,6 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
146
  # fig.savefig('wh.png', dpi=200)
147
 
148
  # Evolve
149
- npr = np.random
150
  f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
151
  pbar = tqdm(range(gen), desc=f'{PREFIX}Evolving anchors with Genetic Algorithm:') # progress bar
152
  for _ in pbar:
 
1
  # YOLOv5 πŸš€ by Ultralytics, GPL-3.0 license
2
  """
3
+ AutoAnchor utils
4
  """
5
 
6
  import random
 
81
  """
82
  from scipy.cluster.vq import kmeans
83
 
84
+ npr = np.random
85
  thr = 1 / thr
86
 
87
  def metric(k, wh): # compute metrics
 
122
  if i:
123
  LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
124
  wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
125
+ # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
126
 
127
  # Kmeans calculation
128
  LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
129
  s = wh.std(0) # sigmas for whitening
130
+ k = kmeans(wh / s, n, iter=30)[0] * s # points
131
+ if len(k) != n: # kmeans may return fewer points than requested if wh is insufficient or too similar
132
+ LOGGER.warning(f'{PREFIX}WARNING: scipy.cluster.vq.kmeans returned only {len(k)} of {n} requested points')
133
+ k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init
134
  wh = torch.tensor(wh, dtype=torch.float32) # filtered
135
  wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
136
  k = print_results(k, verbose=False)
 
148
  # fig.savefig('wh.png', dpi=200)
149
 
150
  # Evolve
 
151
  f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
152
  pbar = tqdm(range(gen), desc=f'{PREFIX}Evolving anchors with Genetic Algorithm:') # progress bar
153
  for _ in pbar: