Upload autoanchor.py
Browse files- utils/autoanchor.py +160 -0
utils/autoanchor.py
ADDED
@@ -0,0 +1,160 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Auto-anchor utils
|
2 |
+
|
3 |
+
import numpy as np
|
4 |
+
import torch
|
5 |
+
import yaml
|
6 |
+
from scipy.cluster.vq import kmeans
|
7 |
+
from tqdm import tqdm
|
8 |
+
|
9 |
+
from utils.general import colorstr
|
10 |
+
|
11 |
+
|
12 |
+
def check_anchor_order(m):
|
13 |
+
# Check anchor order against stride order for YOLO Detect() module m, and correct if necessary
|
14 |
+
a = m.anchor_grid.prod(-1).view(-1) # anchor area
|
15 |
+
da = a[-1] - a[0] # delta a
|
16 |
+
ds = m.stride[-1] - m.stride[0] # delta s
|
17 |
+
if da.sign() != ds.sign(): # same order
|
18 |
+
print('Reversing anchor order')
|
19 |
+
m.anchors[:] = m.anchors.flip(0)
|
20 |
+
m.anchor_grid[:] = m.anchor_grid.flip(0)
|
21 |
+
|
22 |
+
|
23 |
+
def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
24 |
+
# Check anchor fit to data, recompute if necessary
|
25 |
+
prefix = colorstr('autoanchor: ')
|
26 |
+
print(f'\n{prefix}Analyzing anchors... ', end='')
|
27 |
+
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
|
28 |
+
shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
29 |
+
scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
|
30 |
+
wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh
|
31 |
+
|
32 |
+
def metric(k): # compute metric
|
33 |
+
r = wh[:, None] / k[None]
|
34 |
+
x = torch.min(r, 1. / r).min(2)[0] # ratio metric
|
35 |
+
best = x.max(1)[0] # best_x
|
36 |
+
aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold
|
37 |
+
bpr = (best > 1. / thr).float().mean() # best possible recall
|
38 |
+
return bpr, aat
|
39 |
+
|
40 |
+
anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
|
41 |
+
bpr, aat = metric(anchors)
|
42 |
+
print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
|
43 |
+
if bpr < 0.98: # threshold to recompute
|
44 |
+
print('. Attempting to improve anchors, please wait...')
|
45 |
+
na = m.anchor_grid.numel() // 2 # number of anchors
|
46 |
+
try:
|
47 |
+
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
48 |
+
except Exception as e:
|
49 |
+
print(f'{prefix}ERROR: {e}')
|
50 |
+
new_bpr = metric(anchors)[0]
|
51 |
+
if new_bpr > bpr: # replace anchors
|
52 |
+
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
53 |
+
m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
|
54 |
+
check_anchor_order(m)
|
55 |
+
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
56 |
+
print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
57 |
+
else:
|
58 |
+
print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.')
|
59 |
+
print('') # newline
|
60 |
+
|
61 |
+
|
62 |
+
def kmean_anchors(path='./data/coco.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
|
63 |
+
""" Creates kmeans-evolved anchors from training dataset
|
64 |
+
|
65 |
+
Arguments:
|
66 |
+
path: path to dataset *.yaml, or a loaded dataset
|
67 |
+
n: number of anchors
|
68 |
+
img_size: image size used for training
|
69 |
+
thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
|
70 |
+
gen: generations to evolve anchors using genetic algorithm
|
71 |
+
verbose: print all results
|
72 |
+
|
73 |
+
Return:
|
74 |
+
k: kmeans evolved anchors
|
75 |
+
|
76 |
+
Usage:
|
77 |
+
from utils.autoanchor import *; _ = kmean_anchors()
|
78 |
+
"""
|
79 |
+
thr = 1. / thr
|
80 |
+
prefix = colorstr('autoanchor: ')
|
81 |
+
|
82 |
+
def metric(k, wh): # compute metrics
|
83 |
+
r = wh[:, None] / k[None]
|
84 |
+
x = torch.min(r, 1. / r).min(2)[0] # ratio metric
|
85 |
+
# x = wh_iou(wh, torch.tensor(k)) # iou metric
|
86 |
+
return x, x.max(1)[0] # x, best_x
|
87 |
+
|
88 |
+
def anchor_fitness(k): # mutation fitness
|
89 |
+
_, best = metric(torch.tensor(k, dtype=torch.float32), wh)
|
90 |
+
return (best * (best > thr).float()).mean() # fitness
|
91 |
+
|
92 |
+
def print_results(k):
|
93 |
+
k = k[np.argsort(k.prod(1))] # sort small to large
|
94 |
+
x, best = metric(k, wh0)
|
95 |
+
bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
|
96 |
+
print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr')
|
97 |
+
print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, '
|
98 |
+
f'past_thr={x[x > thr].mean():.3f}-mean: ', end='')
|
99 |
+
for i, x in enumerate(k):
|
100 |
+
print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg
|
101 |
+
return k
|
102 |
+
|
103 |
+
if isinstance(path, str): # *.yaml file
|
104 |
+
with open(path) as f:
|
105 |
+
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict
|
106 |
+
from utils.datasets import LoadImagesAndLabels
|
107 |
+
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
108 |
+
else:
|
109 |
+
dataset = path # dataset
|
110 |
+
|
111 |
+
# Get label wh
|
112 |
+
shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
113 |
+
wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh
|
114 |
+
|
115 |
+
# Filter
|
116 |
+
i = (wh0 < 3.0).any(1).sum()
|
117 |
+
if i:
|
118 |
+
print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
|
119 |
+
wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
|
120 |
+
# wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
|
121 |
+
|
122 |
+
# Kmeans calculation
|
123 |
+
print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
|
124 |
+
s = wh.std(0) # sigmas for whitening
|
125 |
+
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
|
126 |
+
assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}')
|
127 |
+
k *= s
|
128 |
+
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
129 |
+
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
130 |
+
k = print_results(k)
|
131 |
+
|
132 |
+
# Plot
|
133 |
+
# k, d = [None] * 20, [None] * 20
|
134 |
+
# for i in tqdm(range(1, 21)):
|
135 |
+
# k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance
|
136 |
+
# fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
|
137 |
+
# ax = ax.ravel()
|
138 |
+
# ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
|
139 |
+
# fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh
|
140 |
+
# ax[0].hist(wh[wh[:, 0]<100, 0],400)
|
141 |
+
# ax[1].hist(wh[wh[:, 1]<100, 1],400)
|
142 |
+
# fig.savefig('wh.png', dpi=200)
|
143 |
+
|
144 |
+
# Evolve
|
145 |
+
npr = np.random
|
146 |
+
f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
|
147 |
+
pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar
|
148 |
+
for _ in pbar:
|
149 |
+
v = np.ones(sh)
|
150 |
+
while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
|
151 |
+
v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
|
152 |
+
kg = (k.copy() * v).clip(min=2.0)
|
153 |
+
fg = anchor_fitness(kg)
|
154 |
+
if fg > f:
|
155 |
+
f, k = fg, kg.copy()
|
156 |
+
pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
|
157 |
+
if verbose:
|
158 |
+
print_results(k)
|
159 |
+
|
160 |
+
return print_results(k)
|