anchors.py

import numpy as np
import torch
import yaml
from scipy.cluster.vq import kmeans
from tqdm import tqdm
import numpy as np
import torch
import yaml
from tqdm import tqdm
import argparse
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torch.utils.tensorboard import SummaryWriter

from utils.loss import CombinedLoss
from utils.config import load_config
from models.model import Model
from datasets.loaders import get_loader

import matplotlib.pyplot as plt

import utils.augmentations as A


# Adapted from YoloV5
def kmean_anchors(loader, n=9, img_size=640, thr=4.0, gen=1000, verbose=True, max_limit=5000):
    """ Creates kmeans-evolved anchors from training dataset

        Arguments:
            path: path to dataset *.yaml, or a loaded dataset
            n: number of anchors
            img_size: image size used for training
            thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
            gen: generations to evolve anchors using genetic algorithm
            verbose: print all results

        Return:
            k: kmeans evolved anchors

        Usage:
            from utils.autoanchor import *; _ = kmean_anchors()
    """
    thr = 1. / thr

    def metric(k, wh):  # compute metrics
        r = wh[:, None] / k[None]
        x = torch.min(r, 1. / r).min(2)[0]  # ratio metric
        # x = wh_iou(wh, torch.tensor(k))  # iou metric
        return x, x.max(1)[0]  # x, best_x

    def anchor_fitness(k):  # mutation fitness
        _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
        return (best * (best > thr).float()).mean()  # fitness

    def print_results(k):
        k = k[np.argsort(k.prod(1))]  # sort small to large
        x, best = metric(k, wh0)
        bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n  # best possible recall, anch > thr
        print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat))
        print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' %
              (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='')
        for i, x in enumerate(k):
            print('%i,%i' % (round(x[0]), round(x[1])), end=',  ' if i < len(k) - 1 else '\n')  # use in *.cfg
        return k

    shapes = []
    labels = []
    opened = 0
    for img, labs in tqdm(loader):
        if opened > max_limit: break
        labels.append(labs[1])
        for i in range(len(labs[1])):
            shapes.append(img.shape[2:])
        opened += 1
    labels = np.vstack(labels)
    shapes = np.array(shapes)
    #if not (labels[:, 1:] <= 1).all():
    #    # normalize label
    #    labels[:, [2, 4]] /= dataset.shapes[0]
    #    labels[:, [1, 3]] /= dataset.shapes[1]
    # Get label wh
    labels = labels[:, 1:] #drop img id
    shapes = img_size * shapes / shapes.max()
    shapes = shapes[:, [1,0]] # switch order w, h instead of h,w
    # wh0 = np.concatenate([l[:, 3:5] * shapes for l in labels])  # wh
    print(labels[:5])
    print(shapes[:5])
    wh0 = labels[:, 3:5] * shapes
    print(wh0[:5])
    # Filter
    i = (wh0 < 3.0).any(1).sum()
    if i:
        print('WARNING: Extremely small objects found. '
              '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0)))
    wh = wh0[(wh0 >= 2.0).any(1)]  # filter > 2 pixels

    wh1 = wh[wh[:,0] > 2*wh[:,1]]
    s1 = wh1.std(0)  # sigmas for whitening
    k1, dist1 = kmeans(wh1 / s1, 3, iter=30) 
    k1 *= s1

    # Kmeans calculation
    print('Running kmeans for %g anchors on %g points...' % (n, len(wh)))
    s = wh.std(0)  # sigmas for whitening
    k, dist = kmeans(wh / s, n, iter=30)  # points, mean distance
    k *= s
    wh = torch.tensor(wh, dtype=torch.float32)  # filtered
    wh0 = torch.tensor(wh0, dtype=torch.float32)  # unfiltered
    wh1 = torch.tensor(wh1, dtype=torch.float32)
    print_results(k1)
    k = print_results(k)

    # Plot
    """
    k, d = [None] * 20, [None] * 20
    for i in tqdm(range(1, 21)):
        k[i-1], d[i-1] = kmeans(wh / s, i)  # points, mean distance
    fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
    ax = ax.ravel()
    ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
    fig, ax = plt.subplots(1, 2, figsize=(14, 7))  # plot wh
    ax[0].hist(wh[wh[:, 0]<100, 0],400)
    ax[1].hist(wh[wh[:, 1]<100, 1],400)
    #plt.show()
    fig.savefig('wh.png', dpi=200)
    """

    # Evolve
    npr = np.random
    f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1  # fitness, generations, mutation prob, sigma
    pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm')  # progress bar
    for _ in pbar:
        v = np.ones(sh)
        while (v == 1).all():  # mutate until a change occurs (prevent duplicates)
            v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
        kg = (k.copy() * v).clip(min=2.0)
        fg = anchor_fitness(kg)
        if fg > f:
            f, k = fg, kg.copy()
            pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f
            if verbose:
                print_results(k)

    return print_results(k)


parser = argparse.ArgumentParser()
parser.add_argument('-cfg', '--config', type=str, help="Path to training config", required=True)
parser.add_argument('-n', '--num_anchors', type=int, help="Number of anchors", required=True)

args = parser.parse_args()

cfg = load_config(args.config)
transforms = A.Compose([
    A.RandomCropToAspect(cfg.img_shape),
    A.Resize(cfg.img_shape)
])

trainloader = get_loader(cfg.dataset, "train", cfg.dataset_dir, 1, transforms=transforms, shuffle=True)

kmean_anchors(trainloader, n = 9, img_size = max(cfg.img_shape), max_limit=10000)