image_processing/brain_segmentation.cpp

/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/
 
#include <arrayfire.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include "../common/progress.h"
 
using namespace af;
 
const float h_sx_kernel[] = {1, 2, 1, 0, 0, 0, -1, -2, -1};
const float h_sy_kernel[] = {-1, 0, 1, -2, 0, 2, -1, 0, 1};
 
// Unused
// const float h_lp_kernel[] = { -0.5f, -1.0f, -0.5f,
//    -1.0f,  6.0f, -1.0f,
//    -0.5f, -1.0f, -0.5f
//};
 
array edges_slice(array x) {
    array ret;
    static array kernelx = array(dim4(3, 3), h_sx_kernel);
    static array kernely = array(dim4(3, 3), h_sy_kernel);
    ret                  = convolve(x, kernelx) + convolve(x, kernely);
    return abs(ret);
}
 
array gauss(array x, float u, float s) {
    double f = 1 / sqrt(2 * af::Pi * s * s);
    array e  = exp(-pow((x - u), 2) / (2 * s * s));
    return f * e;
}
 
array segment_volume(array A, int k) {
    array I1 = A(span, span, k);
 
    float mx = max<float>(I1);
    float mn = min<float>(I1);
 
    float u0 = 0.9 * mx;
    float s0 = (mx - mn) / 2;
 
    float u1 = 1.1 * mn;
    float s1 = (mx - mn) / 2;
 
    array L0  = gauss(I1, u0, s0);
    array L11 = gauss(I1, u1, s1);
    array L10;
    array L12;
    static array kernel = constant(1, 3, 3) / 9;
    static array L11_old;
    static array L12_old;
 
    if (k == 0) {
        L11 = convolve(L11, kernel);
        L10 = L11;
    } else {
        L10 = L11_old;
        L11 = L12_old;
    }
 
    if (k < A.dims(2) - 1) {
        L12 = gauss(A(span, span, k + 1), u1, s1);
        L12 = convolve(L12, kernel);
    } else {
        L12 = L11;
    }
 
    L11_old = L11;
    L12_old = L12;
 
    array L1 = (L10 + L11 + L12) / 3;
    array S  = (L0 > L1);
    return S.as(A.type());
}
 
void brain_seg(bool console) {
    af::Window wnd("Brain Segmentation Demo");
    wnd.setColorMap(AF_COLORMAP_HEAT);
 
    double time_total = 30;  // run for N seconds
 
    array B    = loadImage(ASSETS_DIR "/examples/images/brain.png");
    int slices = 256;
 
    B = moddims(B, dim4(B.dims(0), B.dims(1) / slices, slices));
    af::sync();
 
    int N = 2 * slices - 1;
 
    timer t  = timer::start();
    int iter = 0;
 
    /* loop forward and backward for 100 frames
     * exit if the user presses escape or the animation
     * ends
     */
    for (int i = 0; !wnd.close(); i++) {
        iter++;
 
        int j    = i % N;
        int k    = std::min(j, N - j);
        array Bi = B(span, span, k);
 
        /* process */
        array Si = segment_volume(B, k);
        array Ei = edges_slice(Si);
        array Mi = meanShift(Bi, 10, 10, 5);
 
        /* visualization */
        if (!console) {
            wnd.grid(2, 2);
 
            wnd(0, 0).image(Bi / 255.f, "Input");
            wnd(1, 0).image(Ei, "Edges");
            wnd(0, 1).image(Mi / 255.f, "Meanshift");
            wnd(1, 1).image(Si, "Segmented");
 
            wnd.show();
        } else {
            /* sync the operations so that current
             * iteration comptation finishes
             * */
            af::sync();
        }
 
        /* we have had ran throuh simlation results
         * exit the rendering loop */
        if (!progress(iter, t, time_total)) break;
        if (!(i < 100 * N)) break;
    }
}
 
int main(int argc, char* argv[]) {
    int device   = argc > 1 ? atoi(argv[1]) : 0;
    bool console = argc > 2 ? argv[2][0] == '-' : false;
 
    try {
        af::setDevice(device);
        af::info();
 
        printf("Brain segmentation example\n");
        brain_seg(console);
 
    } catch (af::exception& e) { fprintf(stderr, "%s\n", e.what()); }
 
    return 0;
}