Image Segmentation with Distance Transform and Watershed Algorithm

#include <cv.hpp>
#include <iostream>
 
using namespace std;
using namespace cv;
 
int DELAY_CAPTION = 1500;
char window_name[] = "Watershed Demo";
Mat src;
 
int display_caption(char *caption, Mat input);
 
int main()
{
    src = imread("img/card2.png");
 
    if (!src.data)
        return -1;
        
    display_caption("source Image", src);
 
    for (int x = 0; x < src.rows; x++) {
        for (int y = 0; y < src.cols; y++) {
            if (src.at<Vec3b>(x, y) == Vec3b(255, 255, 255)) {
                src.at<Vec3b>(x, y)[0] = 0;
                src.at<Vec3b>(x, y)[1] = 0;
                src.at<Vec3b>(x, y)[2] = 0;
            }
        }
    }
 
    Mat kernel = (Mat_<float>(3,3) << 1, 1, 1,
                                      1, -8, 1,
                                      1, 1, 1);
 
    Mat imgLaplacian;
    Mat sharp = src;
    filter2D(sharp, imgLaplacian, CV_32F, kernel);
    src.convertTo(sharp, CV_32F);
    Mat imgResult = sharp - imgLaplacian;
 
    imgResult.convertTo(imgResult, CV_8UC3);
    imgLaplacian.convertTo(imgLaplacian, CV_8UC3);
 
    display_caption("New Sharped Image", imgResult);
    display_caption("New Sharped Image Laplacian", imgLaplacian);
 
    src = imgResult;
 
    Mat bw;
    cvtColor(src, bw, CV_BGR2GRAY);
    threshold(bw, bw, 40, 255, CV_THRESH_BINARY | CV_THRESH_OTSU);
 
    Mat dist;
 
    // CV_DIST_L2 = Euclidean distance
    distanceTransform(bw, dist, CV_DIST_L2, 3);
 
    normalize(dist, dist, 0, 1., NORM_MINMAX);
    display_caption("Distance Transform Image", dist);
 
    threshold(dist, dist, .4, 1., CV_THRESH_BINARY);
 
    Mat kernel1 = Mat::ones(3, 3, CV_8UC1);
    dilate(dist, dist, kernel1);
    display_caption("peaks point", dist);
 
    Mat dist_8U;
    dist.convertTo(dist_8U, CV_8U);
    
    vector<vector<Point>> contours;
    findContours(dist_8U, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);
 
    Mat markers = Mat::zeros(dist.size(), CV_32SC1);
 
    for (size_t i = 0; i < contours.size(); i++)
        drawContours(markers, contours, static_cast<int>(i), Scalar::all(static_cast<int>(i)+1));
 
    circle(markers, Point(5, 5), 3, CV_RGB(255, 255, 255), -1);
    display_caption("Marker", markers * 10000);
 
    watershed(src, markers);
 
    Mat mark = Mat::zeros(markers.size(), CV_8UC1);
    markers.convertTo(mark, CV_8UC1);
    bitwise_not(mark, mark);
    
    display_caption("Markers_v2", mark);
 
    vector<Vec3b>colors;
 
    for (size_t i = 0; i < contours.size(); i++)
    {
        int b = theRNG().uniform(0, 255);
        int g = theRNG().uniform(0, 255);
        int r = theRNG().uniform(0, 255);
 
        colors.push_back(Vec3b((uchar)b, (uchar)g, (uchar)r));
    }
 
    Mat dst = Mat::zeros(markers.size(), CV_8UC3);
 
    for (int i = 0; i < markers.rows; i++)
    {
        for (int j = 0; j < markers.cols; j++)
        {
            int index = markers.at<int>(i, j);
            if (index > 0 && index <= static_cast<int>(contours.size()))
                dst.at<Vec3b>(i, j) = colors[index - 1];
            else
                dst.at<Vec3b>(i, j) = Vec3b(0, 0, 0);
        }
    }
 
    display_caption("Final", dst);
 
    waitKey(0); 
    
    return 0;
}
 
int display_caption(char *caption, Mat input)
{
    Mat dst;
    dst = Mat::zeros(input.size(), input.type());
    
    putText(dst, caption, Point(src.cols / 4, src.rows / 2), FONT_HERSHEY_COMPLEX, 1, Scalar(255, 255, 255));
 
    imshow(window_name, dst);
    int c = waitKey(DELAY_CAPTION);
 
    if (c >= 0)
        return -1;
 
    imshow(window_name, input);
    
    c = waitKey(DELAY_CAPTION);
 
    if (c >= 0)
        return -1;
    return 0;
}