examples/machine_learning/kmeans.cpp

#include <iostream>
#include <stdio.h>
#include <arrayfire.h>
#include "ppm_utils.h"

using namespace af;

int k = 3;


// kmeans(input,input,output)
// data:  input,  1D or 2D (range > [0-1])
// k:     input,  # desired means (k > 1)
// means: output, vector of means
void kmeans(array& data, int k, array& means) {

    array datavec = flat(data);
    float minimum = min<float>(datavec);
    datavec = datavec - minimum; // re-center range

    int nbins = max<float>(datavec) + 1;
    array means_vec = array(seq(k)) * nbins / (k + 1); // initial centroids

    printf("min %g\nmax %d\n", minimum, nbins);

    array hist_counts = constant(0,1, nbins);
    array hist = histogram(datavec, nbins);

    array hist_idx = where(hist);    // non-zero histogram bins
    int num_uniq = hist_idx.elements();

    while (1) { // convergence
        array prev_means = means_vec;

        gfor(array i, num_uniq) {    // current classification
            array diffs = abs(hist_idx(i) - means_vec);
            array val, idx;
            min(val, idx, diffs);
            hist_counts(hist_idx(i)) = idx;
        }

        for (int i = 0; i < k; ++i) {// recacluate means
            array m = where(hist_counts == i);
            means_vec(i) = sum(m * hist(m)) / sum(hist(m));
        }

        if (norm<float>(means_vec - prev_means) < 1) { break; }
    }

    means = means_vec + minimum; // re-center to original range
}


// Re-map input data to k mean values.
array get_mask(array& data, array& means_vec, int k) {
    array vol = constant(0,data.dims()[0], data.dims()[1], k);
    gfor(array i, k) { // pixel differences from computed means
        vol(span, span, i) = abs(data - means_vec(i));
    }
    array mask = constant(0,data.dims());
    array idx, val;
    min(val, idx, vol, 2);  // get lables from min differences
    for(int i=0; i<k; ++i){
        mask = mask + means_vec(i) * (idx==i);
    }
    return mask;
}


// kmeans(input,input,output,output)
// data:  input,  1D or 2D (range > [0-1])
// k:     input,  # desired means (k > 1)
// means: output, vector of means
// mask:  output, data 'shifted' to mean values
void kmeans(array& data, int k, array& means, array& mask) {
    kmeans(data, k, means);
    mask = get_mask(data, means, k);
}


// K-Means image recoloring.
// Shifts the hues of an image to the k mean hues.
void kmeans_demo(bool console) {
    printf("k =  %d\n", k);
    array img = load_gray_ppm("../image_processing/len_std.ppm") * 255.f; // [0-255]
    array means, mask;
    kmeans(img, k, means, mask); // get k means and re-colored image
    print(means);
    if (!console) {
        fig("color","gray");
        fig("sub",2,1,1); image(img); fig("title","input");
        fig("sub",2,1,2); image(mask); fig("title","kmeans shift");
        fig("draw");
    }
}


// Entry point.
int main(int argc, char** argv) {

    bool console = false;
    if (argc > 2 || (argc == 2 && argv[1][0] != '-')) {
        printf("usage: kmeans_demo [-]\n");
        return -1;
    } else if (argc == 2 && argv[1][0] == '-') {
        console = true;
    }

    try {
        printf("** ArrayFire K-Means Demo **\n\n");
        kmeans_demo(console);

    } catch (af::exception& e) {
        fprintf(stderr, "%s\n", e.what());
    }

    if (!console) {
        printf("hit [enter]...");
        getchar();
    }
    return 0;
}