diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..7f5bc0f --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,39 @@ +cmake_minimum_required(VERSION 2.6) +project(segmentationtree) +set(PROJECT_NAME segmentation) + +set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/cmake) + +set(CMAKE_C_COMPILER "gcc-4.9") +set(CMAKE_CXX_COMPILER "g++-4.9") + +# Boost +set(Boost_USE_STATIC_LIBS ON) +set(Boost_USE_MULTITHREADED ON) +find_package(Boost COMPONENTS python REQUIRED) + +find_package(OpenCV REQUIRED) +#find_package(PythonLibs REQUIRED) +#find_package(NumPy REQUIRED) + +set(CMAKE_CXX_FLAGS "-std=c++11 -O3") + +#include_directories( +# ${PYTHON_INCLUDE_DIRS} +# ${NUMPY_INCLUDE_DIRS} +#) + +#add_library(segmentation_tree MODULE src/tree.cpp) +#set_target_properties(segmentation_tree PROPERTIES PREFIX "") +#target_link_libraries(segmentation_tree +# ${Boost_LIBRARIES} +# ${OpenCV_LIBRARIES} +# ${PYTHON_LIBRARIES} +#) + +add_executable(segment src/segment.cpp) +target_link_libraries(segment + ${Boost_LIBRARIES} + ${OpenCV_LIBRARIES} + ${PYTHON_LIBRARIES} +) diff --git a/src/edge.h b/src/edge.h new file mode 100644 index 0000000..d77a0eb --- /dev/null +++ b/src/edge.h @@ -0,0 +1,23 @@ +#ifndef EDGE_H_ +#define EDGE_H_ + +class Edge { +public: + Edge(uint32_t a, uint32_t b, float w) : a_(a), b_(b), weight_(w) {} + Edge() {} + + inline float weight() { return weight_; } + inline uint32_t first() { return a_; } + inline uint32_t second() { return b_; } + + inline void setFirst(int a) { a_ = a; } + inline void setSecond(int b) { b_ = b; } + inline void setWeight(float w) { weight_ = w; } + +private: + uint32_t a_; + uint32_t b_; + float weight_; +}; + +#endif diff --git a/src/segment.cpp b/src/segment.cpp new file mode 100644 index 0000000..fcb3198 --- /dev/null +++ b/src/segment.cpp @@ -0,0 +1,199 @@ +#include +#include +#include +#include + +#include +#include +#include + +#include "edge.h" + +bool sort(Edge a, Edge b) { return a.weight() < b.weight(); } + +// TODO: compare with struct in a vector? +// cluster data: +// [0] = id +// [1] = threshold +// [2] = size of cluster +// [3] = rank (number of times merged with other cluster) + +/// Join two clusters +void join(cv::Vec4f & c1, cv::Vec4f & c2, int k, float weight) { + if (c1[3] > c2[3]) { + c2[0] = c1[0]; + c1[2] += c2[2]; + c1[1] = weight + k / c1[2]; + } else { + c1[0] = c2[0]; + c2[2] += c1[2]; + if (c1[3] == c2[3]) + c2[3]++; + c2[1] = weight + k / c2[2]; + } +} + +/// Difference between two pixels p1, p2 in image +float diff(cv::Mat image, int p1, int p2) { + cv::Vec3b d = image.at(p1) - image.at(p2); + return sqrt(d[0]*d[0] + d[1]*d[1] + d[2]*d[2]); +} + + +/// Traverses the path in clusters to find the end cluster for id +cv::Vec4f& find(cv::Mat & clusters, int id) { + cv::Vec4f & c = clusters.at(id); + cv::Vec4f & c_ = clusters.at(c[0]); + while (c[0] != c_[0]) { + c = c_; + c_ = clusters.at(c[0]); + } + clusters.at(id)[0] = c[0]; + return c; +} + +/// Main segment method +cv::Mat segment(cv::Mat image, double sigma, int k, int min_size) { + cv::GaussianBlur(image, image, cv::Size(3, 3), sigma); + + //cv::Mat clusters(image.size(), CV_32SC1); + cv::Mat clusters = cv::Mat::zeros(image.size(), CV_32FC4); + //cv::Mat sizes(image.size(), CV_32SC1); + + std::vector edges; + edges.resize(image.rows*image.cols*4); + + // add edges + int index = 0; + for (int i = 1; i < image.rows - 1; i++) { + for (int j = 0; j < image.cols - 1; j++) { + int v1_ = i * image.cols + j; + cv::Vec4f & c = clusters.at(i, j); + c[0] = v1_; + c[1] = k; + + // right + //if (j < image.cols - 1) { + int v2_ = v1_ + 1; + //int v2_ = i * image.cols + (j+1); + Edge & e = edges[index++]; + e.setFirst(v1_); + e.setSecond(v2_); + e.setWeight(diff(image, v1_, v2_)); + //} + // down + //if (i < image.rows - 1) { + v2_ = v1_ + image.cols; + //int v2_ = (i+1) * image.cols + j; + e = edges[index++]; + e.setFirst(v1_); + e.setSecond(v2_); + e.setWeight(diff(image, v1_, v2_)); + //} + // down right + //if (i < image.rows - 1 && j < image.cols - 1) { + v2_ = v1_ + image.cols + 1; + //int v2_ = (i+1) * image.cols + (j+1); + e = edges[index++]; + e.setFirst(v1_); + e.setSecond(v2_); + e.setWeight(diff(image, v1_, v2_)); + //} + // up right + //if (i > 0 && j < image.cols - 1) { + v2_ = v1_ - image.cols + 1; + //int v2_ = (i-1) * image.cols + (j+1); + e = edges[index++]; + e.setFirst(v1_); + e.setSecond(v2_); + e.setWeight(diff(image, v1_, v2_)); + //} + } + } + + // add cluster data for skipped borders + // last column + for (int i = 0; i < image.rows; i++) { + cv::Vec4f & c = clusters.at(i, image.cols-1); + c[0] = i * image.cols + image.cols - 1; + c[1] = k; + } + // first row + for (int j = 0; j < image.cols; j++) { + cv::Vec4f & c = clusters.at(0, j); + c[0] = j; + c[1] = k; + } + // last row + for (int j = 0; j < image.cols; j++) { + cv::Vec4f & c = clusters.at(image.rows-1, j); + c[0] = (image.rows-1) * image.cols + j; + c[1] = k; + } + + //std::clock_t begin = std::clock(); + std::sort(edges.begin(), edges.begin() + index, sort); + //std::clock_t end = std::clock(); + //double elapsed_secs = double(end - begin) / CLOCKS_PER_SEC; + //std::cout << "Times passed in seconds: " << elapsed_secs << std::endl; + + // merge the clusters + for (auto e: edges) { + cv::Vec4f & c1 = find(clusters, e.first()); + cv::Vec4f & c2 = find(clusters, e.second()); + if (c1[0] != c2[0] && + e.weight() <= c1[1] && + e.weight() <= c2[1]) { + join(c1, c2, k, e.weight()); + } + } + + // choose random colors + cv::Mat colors(image.size(), CV_8UC3); + for (int i = 0; i < colors.total(); i++) { + colors.at(i) = {uint8_t(rand() % 256), uint8_t(rand() % 256), uint8_t(rand() % 256)}; + } + + //std::map colormap; + + // color the segmentation + cv::Mat segmentation = cv::Mat::zeros(image.size(), CV_8UC3); + for (int i = 1; i < image.rows - 1; i++) { + for (int j = 0; j < image.cols - 1; j++) { + int index = i * image.cols + j; + cv::Vec4f & c = find(clusters, index); + segmentation.at(index) = colors.at(c[0]); + } + } + /*for (auto v: vertices) { + auto color = colormap.find(v.cluster()); + if (color == colormap.end()) { + colormap[v.cluster()] = cv::Vec3b(rand() % 256, rand() % 256, rand() % 256); + } + segmentation.at(v.row(), v.col()) = colormap[v.cluster()]; + }*/ + + return segmentation; +} + +int main(int argc, char * argv[]) { + if (argc != 2) { + std::cout << "Provide image as first argument." << std::endl; + return 0; + } + + cv::namedWindow("Image", cv::WINDOW_NORMAL); + cv::Mat image = cv::imread(argv[1]); + + cv::Mat segmentation; + std::clock_t begin = std::clock(); + for (int i = 0; i < 10; i++) + segmentation = segment(image, 0.8, 300, 20); + std::clock_t end = std::clock(); + double elapsed_secs = double(end - begin) / CLOCKS_PER_SEC / 10; + std::cout << "Times passed in seconds: " << elapsed_secs << std::endl; + + cv::imshow("Image", segmentation); + cv::waitKey(); + return 0; +}