<repo_name>satojkovic/fm-detector<file_sep>/README.md
# fm-detector
under construction.
<file_sep>/main.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Usage:
    main.py <image_file>
    main.py -h | --help
Options:
    -h --help  show help message
"""
import os
import numpy as np
import cv2
from collections import Counter
from math import sqrt
import sys
from docopt import docopt


rtoh = lambda rgb: '#%s' % ''.join(('%02x' % p for p in rgb))


def get_colors(img, w, h):
    pixels = [(img[y, x][0], img[y, x][1], img[y, x][2])
              for x in range(w) for y in range(h)]
    return Counter(pixels)


def get_points(img):
    points = []
    h, w, ch = img.shape
    for color, count in get_colors(img, w, h).items():
        points.append([np.array(color), count])
    return np.array(points)


def initialize(data, k):
    # choose one data point as an initial cluster center
    center = [data[0][0]]

    for i in range(1, k):
        # calculate distances
        dist = np.array([np.min([np.inner(c-d[0], c-d[0]) for c in center])
                         for d in data])

        # probability distribution
        probs = dist / dist.sum()
        cumprobs = np.cumsum(probs)

        # select next cluster center
        r = np.random.random()
        for j, p in enumerate(cumprobs):
            if r < p:
                i = j
                break
        center.append(data[i][0])
    return np.array(center)


def calculate_center(data_i):
    val = 0.0
    cnt_i = len(data_i)
    for d in data_i:
        val += np.array(d[0])
    return (val / cnt_i)


def euclid_dist(center, data_i):
    center = center.astype(np.int)
    data_i = data_i.astype(np.int)
    return sqrt(sum([
        (center[n] - data_i[n]) ** 2 for n in range(len(center))
    ]))


def assign_label(center, data_i):
    dist = []
    for k in range(len(center)):
        dist.append(euclid_dist(center[k], data_i[0]))
    return np.argsort(np.ravel(dist))[0]


def kmeanspp(data, k, max_iter=300):
    # calculate an initial cluster center
    center = initialize(data, k)

    # assign the cluster labels
    labels = np.zeros(len(data)).astype(np.int)
    for i, d in enumerate(data):
        labels[i] = assign_label(center, d)

    # kmeans clustering
    iter = 0
    while 1:
        print iter, 'th iteration'

        # calculate the cluster centers
        for i in range(k):
            center[i] = calculate_center(data[labels == i])

        # assign new labels
        new_labels = np.zeros(len(data)).astype(np.int)
        for i, d in enumerate(data):
            new_labels[i] = assign_label(center, d)

        if np.array_equal(labels, new_labels) or iter > max_iter:
            break
        else:
            labels = new_labels

        iter += 1

    return new_labels, center


def dominant_colors(filename, k=5):
    img = cv2.imread(filename)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    points = get_points(img)
    labels, center = kmeanspp(points, k)
    rgbs = [map(np.int, c) for c in center]
    return map(rtoh, rgbs)


def main():
    args = docopt(__doc__)
    img_file = args['<image_file>']
    if not os.path.exists(img_file):
        print 'Not found:', img_file
        sys.exit(-1)

    # dominant colors
    print dominant_colors(img_file, k=5)


if __name__ == '__main__':
    main()
