Algorithm should work now

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)

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

src/edge.h

@@ -3,17 +3,19 @@
 
 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_; }
 
 	inline void setFirst(int a) { a_ = a; }
 	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_;

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<Cluster> & 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<Cluster> & 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<Cluster> segment_graph(int num_vertices, int num_edges, std::vector<Edge> & edges, float c, int & num) { 
+	// sort edges by weight
+	std::sort(edges.begin(), edges.begin() + num_edges);
+
+	// make a disjoint-set forest
+	std::vector<Cluster> 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

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<cv::Vec3b>(p1) - image.at<cv::Vec3b>(p2);
+	cv::Vec3f a = image.at<cv::Vec3f>(p1);
+	cv::Vec3f b = image.at<cv::Vec3f>(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<Edge> 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<Cluster> 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<cv::Vec3b> 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<cv::Vec3b>(y, x) = colors[comp];
+		}
+	}  
+
+	return output;
+}
+
+#endif