#!/usr/bin/env python
# vim:tabstop=4 shiftwidth=4 tw=79:

'''
SVM, Random forest and KNearest digit recognition.
Modified from the OpenCV example.

Sample loads a dataset of handwritten digits from '../data/digits.png'.  Then
it trains a Random Forest, SVM and KNearest classifiers on it and evaluates
their accuracy.

Following preprocessing is applied to the dataset:
  - Moment-based image deskew (see deskew())
  - Digit images are split into 4 10x10 cells and 16-bin histogram of oriented
    gradients is computed for each cell
  - Transform histograms to space with Hellinger metric (see [1] (RootSIFT))

Or in the "simple setting":
  - Only moment-based image deskew (see deskew())

[1] R. Arandjelovic, A. Zisserman
    "Three things everyone should know to improve object retrieval"
    http://www.robots.ox.ac.uk/~vgg/publications/2012/Arandjelovic12/arandjelovic12.pdf

Usage:
     digits.py
'''

import cv2

import numpy as np
from numpy.linalg import norm

# local modules
from common import mosaic


SZ = 20  # size of each digit is SZ x SZ
CLASS_N = 10
DIGITS_FN = 'digits.png'
SIMPLE = True  # Use simple preprocessing or HOG features (for SVM)


def split2d(img, cell_size, flatten=True):
    h, w = img.shape[:2]
    sx, sy = cell_size
    cells = [np.hsplit(row, w//sx) for row in np.vsplit(img, h//sy)]
    cells = np.array(cells)
    if flatten:
        cells = cells.reshape(-1, sy, sx)
    return cells


def load_digits(fn):
    print 'loading "%s" ...' % fn
    digits_img = cv2.imread(fn, 0)
    digits = split2d(digits_img, (SZ, SZ))
    labels = np.repeat(np.arange(CLASS_N), len(digits)/CLASS_N)
    return digits, labels


def deskew(img):
    m = cv2.moments(img)
    if abs(m['mu02']) < 1e-2:
        return img.copy()
    skew = m['mu11']/m['mu02']
    M = np.float32([[1, skew, -0.5*SZ*skew], [0, 1, 0]])
    img = cv2.warpAffine(img, M, (SZ, SZ),
                         flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
    return img


class StatModel(object):
    def load(self, fn):
        self.model.load(fn)

    def save(self, fn):
        self.model.save(fn)


class KNearest(StatModel):
    def __init__(self, k=3):
        self.k = k
        self.model = cv2.KNearest()

    def train(self, samples, responses):
        self.model = cv2.KNearest()
        self.model.train(samples, responses)

    def predict(self, samples):
        retval, results, neigh_resp, dists = self.model.find_nearest(samples,
                                                                     self.k)
        return results.ravel()


class SVM(StatModel):
    def __init__(self, kernel_type=cv2.SVM_RBF, C=1, gamma=0.5):
        self.params = dict(kernel_type=kernel_type,
                           svm_type=cv2.SVM_C_SVC,
                           C=C,
                           gamma=gamma)
        self.model = cv2.SVM()

    def train(self, samples, responses):
        self.model = cv2.SVM()
        self.model.train(samples, responses, params=self.params)

    def train_auto(self, samples, responses):
        self.model = cv2.SVM()
        self.model.train_auto(samples, responses, None, None,
                              params=self.params)

    def predict(self, samples):
        return self.model.predict_all(samples).ravel()


class RForest(StatModel):
    def __init__(self):
        self.params = dict(max_depth=20)
        self.model = cv2.RTrees()

    def train(self, samples, responses):
        self.model = cv2.RTrees()
        self.model.train(samples, cv2.CV_ROW_SAMPLE, responses,
                         params=self.params)

    def predict(self, samples):
        predictions = map(self.model.predict, samples)
        return predictions


def evaluate_model(model, digits, samples, labels):
    resp = model.predict(samples)
    err = (labels != resp).mean()
    print 'error: %.2f %%' % (err*100)

    confusion = np.zeros((10, 10), np.int32)
    for i, j in zip(labels, resp):
        confusion[i, j] += 1
    print 'confusion matrix:'
    print confusion
    print

    vis = []
    for img, flag in zip(digits, resp == labels):
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        if not flag:
            img[..., :2] = 0
        vis.append(img)
    return mosaic(25, vis)


def preprocess_simple(digits):
    return np.float32(digits).reshape(-1, SZ*SZ) / 255.0


def preprocess_hog(digits):
    samples = []
    for img in digits:
        gx = cv2.Sobel(img, cv2.CV_32F, 1, 0)
        gy = cv2.Sobel(img, cv2.CV_32F, 0, 1)
        mag, ang = cv2.cartToPolar(gx, gy)
        bin_n = 16
        bin = np.int32(bin_n*ang/(2*np.pi))
        bin_cells = bin[:10, :10], bin[10:, :10], bin[:10, 10:], bin[10:, 10:]
        mag_cells = mag[:10, :10], mag[10:, :10], mag[:10, 10:], mag[10:, 10:]
        hists = [np.bincount(b.ravel(), m.ravel(), bin_n)
                 for b, m in zip(bin_cells, mag_cells)]
        hist = np.hstack(hists)

        # transform to Hellinger kernel
        eps = 1e-7
        hist /= hist.sum() + eps
        hist = np.sqrt(hist)
        hist /= norm(hist) + eps

        samples.append(hist)
    return np.float32(samples)


if __name__ == '__main__':
    print __doc__

    digits, labels = load_digits(DIGITS_FN)

    print 'preprocessing...'
    # shuffle digits
    rand = np.random.RandomState(321)
    shuffle = rand.permutation(len(digits))
    digits, labels = digits[shuffle], labels[shuffle]

    digits2 = map(deskew, digits)
    if SIMPLE:
        samples = preprocess_simple(digits2)
    else:
        samples = preprocess_hog(digits2)

    train_n = int(0.9*len(samples))
    cv2.imshow('test set', mosaic(25, digits[train_n:]))
    digits_train, digits_test = np.split(digits2, [train_n])
    samples_train, samples_test = np.split(samples, [train_n])
    labels_train, labels_test = np.split(labels, [train_n])

    print 'training Random Forest...'

    model = RForest()
    model.train(samples_train, labels_train)
    vis = evaluate_model(model, digits_test, samples_test, labels_test)
    cv2.imshow('Random Forest test', vis)

    print 'training KNearest...'
    model = KNearest(k=4)
    model.train(samples_train, labels_train)
    vis = evaluate_model(model, digits_test, samples_test, labels_test)
    cv2.imshow('KNearest test', vis)

    print 'training SVM...'

    if SIMPLE:
        model = SVM(kernel_type=cv2.SVM_LINEAR, C=0.1)
        model.train(samples_train, labels_train)
    else:
        model = SVM(kernel_type=cv2.SVM_RBF, C=2.67, gamma=5.383)
        model.train(samples_train, labels_train)

    vis = evaluate_model(model, digits_test, samples_test, labels_test)
    cv2.imshow('SVM test', vis)
    print 'saving SVM as "digits_svm.dat"...'
    model.save('digits_svm.dat')

    cv2.waitKey(0)
