From 3c1782586a39b8b1ec60e27a09c94c73f50f1474 Mon Sep 17 00:00:00 2001 From: Hans Gaiser Date: Tue, 6 May 2014 22:04:44 +0200 Subject: [PATCH] Algorithm should work now --- CMakeLists.txt | 4 +- src/cluster.h | 28 ++++++ src/edge.h | 6 +- src/segment-graph.h | 76 +++++++++++++++ src/segment-image.h | 116 ++++++++++++++++++++++ src/segment.cpp | 232 +++++++------------------------------------- 6 files changed, 259 insertions(+), 203 deletions(-) create mode 100644 src/cluster.h create mode 100644 src/segment-graph.h create mode 100644 src/segment-image.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 7f5bc0f..198a4bd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -4,8 +4,8 @@ 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") +#set(CMAKE_C_COMPILER "gcc-4.9") +#set(CMAKE_CXX_COMPILER "g++-4.9") # Boost set(Boost_USE_STATIC_LIBS ON) diff --git a/src/cluster.h b/src/cluster.h new file mode 100644 index 0000000..df7fe53 --- /dev/null +++ b/src/cluster.h @@ -0,0 +1,28 @@ +#ifndef CLUSTER_H_ +#define CLUSTER_H_ + +class Cluster { +public: + Cluster() { + setSize(1); + setRank(0); + } + + inline float threshold() const { return threshold_; } + inline uint32_t size() const { return size_; } + inline uint32_t rank() const { return rank_; } + inline uint32_t id() const { return id_; } + + inline void setThreshold(float t) { threshold_ = t; } + inline void setSize(int s) { size_ = s; } + inline void setRank(int r) { rank_ = r; } + inline void setId(int id) { id_ = id; } + +private: + float threshold_; + uint32_t size_; + uint32_t rank_; + uint32_t id_; +}; + +#endif diff --git a/src/edge.h b/src/edge.h index d77a0eb..6f89823 100644 --- a/src/edge.h +++ b/src/edge.h @@ -3,10 +3,9 @@ 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 float weight() const { return weight_; } inline uint32_t first() { return a_; } inline uint32_t second() { return b_; } @@ -14,6 +13,9 @@ class Edge { inline void setSecond(int b) { b_ = b; } inline void setWeight(float w) { weight_ = w; } + bool operator<(const Edge & b) const { + return weight() < b.weight(); + } private: uint32_t a_; uint32_t b_; diff --git a/src/segment-graph.h b/src/segment-graph.h new file mode 100644 index 0000000..6c39d67 --- /dev/null +++ b/src/segment-graph.h @@ -0,0 +1,76 @@ +#ifndef SEGMENT_GRAPH +#define SEGMENT_GRAPH + +#include "edge.h" +#include "cluster.h" + +// threshold function +#define THRESHOLD(size, c) (c/size) + +int find(std::vector & clusters, int x) { + int y = x; + while (y != clusters[y].id()) { + y = clusters[y].id(); + } + clusters[x].setId(y); + return y; +} + +void join(std::vector & clusters, int x, int y) { + if (clusters[x].rank() > clusters[y].rank()) { + clusters[y].setId(x); + clusters[x].setSize(clusters[x].size() + clusters[y].size()); + } else { + clusters[x].setId(y); + clusters[y].setSize(clusters[y].size() + clusters[x].size()); + if (clusters[x].rank() == clusters[y].rank()) + clusters[y].setRank(clusters[y].rank() + 1); + } +} + +/* + * Segment a graph + * + * Returns a disjoint-set forest representing the segmentation. + * + * num_vertices: number of vertices in graph. + * num_edges: number of edges in graph + * edges: array of edges. + * c: constant for treshold function. + */ +std::vector segment_graph(int num_vertices, int num_edges, std::vector & edges, float c, int & num) { + // sort edges by weight + std::sort(edges.begin(), edges.begin() + num_edges); + + // make a disjoint-set forest + std::vector clusters; + clusters.resize(num_vertices); + + // init thresholds + for (int i = 0; i < num_vertices; i++) { + clusters[i].setId(i); + clusters[i].setThreshold(THRESHOLD(1,c)); + } + + // for each edge, in non-decreasing weight order... + for (int i = 0; i < num_edges; i++) { + Edge & edge = edges[i]; + + // components conected by this edge + int a = find(clusters, edge.first()); + int b = find(clusters, edge.second()); + if (a != b) { + if ((edge.weight() <= clusters[a].threshold()) && + (edge.weight() <= clusters[b].threshold())) { + join(clusters, a, b); + a = find(clusters, a); + clusters[a].setThreshold(edge.weight() + THRESHOLD(clusters[a].size(), c)); + num--; + } + } + } + + return clusters; +} + +#endif diff --git a/src/segment-image.h b/src/segment-image.h new file mode 100644 index 0000000..b804924 --- /dev/null +++ b/src/segment-image.h @@ -0,0 +1,116 @@ +#ifndef SEGMENT_IMAGE +#define SEGMENT_IMAGE + +#include "segment-graph.h" + +#define square(x) ((x)*(x)) + +// random color +cv::Vec3b random_rgb(){ + cv::Vec3b c; + double r; + c[0] = (uchar)random(); + c[1] = (uchar)random(); + c[2] = (uchar)random(); + return c; +} + +/// Difference between two pixels p1, p2 in image +static inline float diff(cv::Mat image, int p1, int p2) { + //cv::Vec3b d = image.at(p1) - image.at(p2); + cv::Vec3f a = image.at(p1); + cv::Vec3f b = image.at(p2); + return sqrt(square(a[0] - b[0]) + square(a[1] - b[1]) + square(a[2] - b[2])); +} + +/* + * Segment an image + * + * Returns a color image representing the segmentation. + * + * im: image to segment. + * sigma: to smooth the image. + * c: constant for treshold function. + * min_size: minimum component size (enforced by post-processing stage). + * num_ccs: number of connected components in the segmentation. + */ +cv::Mat segment_image(cv::Mat image, float sigma, float c, int min_size, int & num_ccs) { + image.convertTo(image, CV_32FC3); + cv::GaussianBlur(image, image, cv::Size(5, 5), sigma); + //image = smooth(image, sigma); + + // build graph + std::vector edges; + edges.resize(image.total() * 4); + int index = 0; + for (int i = 0; i < image.rows; i++) { + for (int j = 0; j < image.cols; j++) { + int v1 = i * image.cols + j; + + if (j < image.cols - 1) { + int v2 = i * image.cols + (j+1); + Edge & e = edges[index++]; + e.setFirst(v1); + e.setSecond(v2); + e.setWeight(diff(image, v1, v2)); + } + if (i < image.rows - 1) { + int v2 = (i+1) * image.cols + j; + Edge & e = edges[index++]; + e.setFirst(v1); + e.setSecond(v2); + e.setWeight(diff(image, v1, v2)); + } + if (i < image.rows - 1 && j < image.cols - 1) { + int v2 = (i+1) * image.cols + (j+1); + Edge & e = edges[index++]; + e.setFirst(v1); + e.setSecond(v2); + e.setWeight(diff(image, v1, v2)); + } + if (i > 0 && j < image.cols - 1) { + int v2 = (i-1) * image.cols + (j+1); + Edge & e = edges[index++]; + e.setFirst(v1); + e.setSecond(v2); + e.setWeight(diff(image, v1, v2)); + } + } + } + + num_ccs = image.total(); + + // segment + std::vector clusters = segment_graph(image.total(), index, edges, c, num_ccs); + + // post process small components + if (min_size) { + for (int i = 0; i < index; i++) { + int a = find(clusters, edges[i].first()); + int b = find(clusters, edges[i].second()); + if ((a != b) && ((clusters[a].size() < min_size) || (clusters[b].size() < min_size))) { + join(clusters, a, b); + num_ccs--; + } + } + } + + cv::Mat output(image.size(), CV_8UC3); + + // pick random colors for each component + std::vector colors; + colors.resize(image.total()); + for (int i = 0; i < image.total(); i++) + colors[i] = random_rgb(); + + for (int y = 0; y < image.rows; y++) { + for (int x = 0; x < image.cols; x++) { + int comp = find(clusters, y * image.cols + x); + output.at(y, x) = colors[comp]; + } + } + + return output; +} + +#endif diff --git a/src/segment.cpp b/src/segment.cpp index fcb3198..16dd886 100644 --- a/src/segment.cpp +++ b/src/segment.cpp @@ -1,199 +1,33 @@ -#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; -} +#include +#include +#include +#include +#include "segment-image.h" +#include + +int main(int argc, char **argv) { + if (argc != 6) { + std::cerr << "usage: " << argv[0] << " sigma k min input output" << std::endl; + return 1; + } + + float sigma = atof(argv[1]); + float k = atof(argv[2]); + int min_size = atoi(argv[3]); + + cv::Mat input = cv::imread(argv[4]); + int num_ccs; + + std::cout << "Processing..." << std::endl; + + std::clock_t begin = std::clock(); + cv::Mat seg = segment_image(input, sigma, k, min_size, num_ccs); + std::clock_t end = std::clock(); + double elapsed_secs = double(end - begin) / CLOCKS_PER_SEC; + cv::imwrite(argv[5], seg); + std::cout << "Times passed in seconds: " << elapsed_secs << std::endl; + + std::cout << "Got " << num_ccs << " components." << std::endl; + return 0; +} +