refactored mssegmentation.cpp

modules/gpu/src/mssegmentation.cpp

@@ -40,38 +40,13 @@
 //
 //M*/
 
-#include <time.h>
-#include <vector>
 #include "precomp.hpp"
 
 #if !defined(HAVE_CUDA)
 
-namespace cv
-{
-namespace gpu
-{
-
-void meanShiftSegmentation(const GpuMat&, Mat&, int, int, int, TermCriteria) { throw_nogpu(); }
-
-} // namespace gpu
-} // namespace cv
-
-#else
-
-//#define _MSSEGMENTATION_DBG
+void cv::gpu::meanShiftSegmentation(const GpuMat&, Mat&, int, int, int, TermCriteria) { throw_nogpu(); }
 
-#ifdef _MSSEGMENTATION_DBG
-#include <iostream>
-#define LOG(s) std::cout << (s) << std::endl
-#define LOG2(s1, s2) std::cout << (s1) << (s2) << std::endl
-#define DBG(code) code
 #else
-#define LOG(s1)
-#define LOG2(s1, s2)
-#define DBG(code)
-#endif
-
-#define PIX(y, x) ((y) * ncols + (x))
 
 using namespace std;
 
@@ -87,13 +62,12 @@ class DjSets
 {
 public:
     DjSets(int n);
-    ~DjSets();
-    int find(int elem) const;
+    int find(int elem);
     int merge(int set1, int set2);
 
-    int* parent;
-    int* rank;
-    int* size;
+    vector<int> parent;
+    vector<int> rank;
+    vector<int> size;
 private:
     DjSets(const DjSets&) {}
     DjSets operator =(const DjSets&) {}
@@ -118,12 +92,11 @@ public:
     typedef GraphEdge<T> Edge;
 
     Graph(int numv, int nume_max);
-    ~Graph();
 
     void addEdge(int from, int to, const T& val=T());
 
-    int* start;
-    Edge* edges;
+    vector<int> start;
+    vector<Edge> edges;
 
     int numv;
     int nume_max;
@@ -152,47 +125,27 @@ struct SegmLink
     SegmLink() {}
     SegmLink(int from, int to, const SegmLinkVal& val) 
         : from(from), to(to), val(val) {}
+    bool operator <(const SegmLink& other) const 
+    {
+        return val < other.val;
+    }
     int from;
     int to;
     SegmLinkVal val;
 };
 
-
-struct SegmLinkCmp
-{
-    bool operator ()(const SegmLink& lhs, const SegmLink& rhs) const
-    {
-        return lhs.val < rhs.val;
-    }
-};
-
 //
 // Implementation
 //
 
-DjSets::DjSets(int n)
+DjSets::DjSets(int n) : parent(n), rank(n, 0), size(n, 1)
 {
-    parent = new int[n];
-    rank = new int[n];
-    size = new int[n];   
     for (int i = 0; i < n; ++i)
-    {
         parent[i] = i;
-        rank[i] = 0;
-        size[i] = 1;
-    }
 }
 
 
-DjSets::~DjSets()
-{
-    delete[] parent;
-    delete[] rank;
-    delete[] size;
-}
-
-
-inline int DjSets::find(int elem) const
+inline int DjSets::find(int elem)
 {
     int set = elem;
     while (set != parent[set])
@@ -229,50 +182,50 @@ inline int DjSets::merge(int set1, int set2)
 
 
 template <typename T>
-Graph<T>::Graph(int numv, int nume_max)
+Graph<T>::Graph(int numv, int nume_max) : start(numv, -1), edges(nume_max)
 {
     this->numv = numv;
     this->nume_max = nume_max;
-    start = new int[numv];
-    for (int i = 0; i < numv; ++i)
-        start[i] = -1;
-    edges = new Edge[nume_max];
     nume = 0;
 }
 
 
 template <typename T>
-Graph<T>::~Graph()
+inline void Graph<T>::addEdge(int from, int to, const T& val)
 {
-    delete[] start;
-    delete[] edges;
+    edges[nume] = Edge(to, start[from], val);
+    start[from] = nume;
+    nume++;
 }
 
 
-template <typename T>
-inline void Graph<T>::addEdge(int from, int to, const T& val)
+inline int pix(int y, int x, int ncols) 
 {
-    Edge* edge = edges + nume;
-    new (edge) SegmLink(to, start[from], val);
-    start[from] = nume;
-    nume++;
+    return y * ncols + x;
 }
 
 
-inline int sqr(int x)
+inline int sqr(int x) 
 {
     return x * x;
 }
 
-} // anonymous namespace
 
-
-namespace cv
+inline int dist2(const cv::Vec4b& lhs, const cv::Vec4b& rhs) 
 {
-namespace gpu
+    return sqr(lhs[0] - rhs[0]) + sqr(lhs[1] - rhs[1]) + sqr(lhs[2] - rhs[2]);
+}
+
+
+inline int dist2(const cv::Vec2s& lhs, const cv::Vec2s& rhs) 
 {
+    return sqr(lhs[0] - rhs[0]) + sqr(lhs[1] - rhs[1]);
+}
+
+} // anonymous namespace
+
 
-void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int minsize, TermCriteria criteria)
+void cv::gpu::meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int minsize, TermCriteria criteria)
 {
     CV_Assert(src.type() == CV_8UC4);
     const int nrows = src.rows;
@@ -280,37 +233,28 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
     const int hr = sr;
     const int hsp = sp;
 
-    DBG(clock_t start = clock());
-
     // Perform mean shift procedure and obtain region and spatial maps
     GpuMat h_rmap, h_spmap;
     meanShiftProc(src, h_rmap, h_spmap, sp, sr, criteria);
     Mat rmap = h_rmap;
     Mat spmap = h_spmap;
 
-    LOG2("meanshift:", clock() - start);
-    DBG(start = clock());
-
     Graph<SegmLinkVal> g(nrows * ncols, 4 * (nrows - 1) * (ncols - 1)
                                         + (nrows - 1) + (ncols - 1));
 
-    LOG2("ragalloc:", clock() - start);
-    DBG(start = clock());
-
     // Make region adjacent graph from image
-    // TODO: SSE?
     Vec4b r1;
     Vec4b r2[4];
-    Point_<short> sp1;
-    Point_<short> sp2[4];
+    Vec2s sp1;
+    Vec2s sp2[4];
     int dr[4];
     int dsp[4];
     for (int y = 0; y < nrows - 1; ++y)
     {
         Vec4b* ry = rmap.ptr<Vec4b>(y);
         Vec4b* ryp = rmap.ptr<Vec4b>(y + 1);
-        Point_<short>* spy = spmap.ptr<Point_<short> >(y);
-        Point_<short>* spyp = spmap.ptr<Point_<short> >(y + 1);
+        Vec2s* spy = spmap.ptr<Vec2s>(y);
+        Vec2s* spyp = spmap.ptr<Vec2s>(y + 1);
         for (int x = 0; x < ncols - 1; ++x)
         {
             r1 = ry[x];
@@ -326,54 +270,48 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
             sp2[2] = spyp[x + 1];
             sp2[3] = spyp[x];
 
-            dr[0] = sqr(r1[0] - r2[0][0]) + sqr(r1[1] - r2[0][1]) + sqr(r1[2] - r2[0][2]);
-            dr[1] = sqr(r1[0] - r2[1][0]) + sqr(r1[1] - r2[1][1]) + sqr(r1[2] - r2[1][2]);
-            dr[2] = sqr(r1[0] - r2[2][0]) + sqr(r1[1] - r2[2][1]) + sqr(r1[2] - r2[2][2]);
-            dsp[0] = sqr(sp1.x - sp2[0].x) + sqr(sp1.y - sp2[0].y);
-            dsp[1] = sqr(sp1.x - sp2[1].x) + sqr(sp1.y - sp2[1].y);
-            dsp[2] = sqr(sp1.x - sp2[2].x) + sqr(sp1.y - sp2[2].y); 
+            dr[0] = dist2(r1, r2[0]);
+            dr[1] = dist2(r1, r2[1]);
+            dr[2] = dist2(r1, r2[2]);
+            dsp[0] = dist2(sp1, sp2[0]);
+            dsp[1] = dist2(sp1, sp2[1]);
+            dsp[2] = dist2(sp1, sp2[2]);
 
             r1 = ry[x + 1];
             sp1 = spy[x + 1];
 
-            dr[3] = sqr(r1[0] - r2[3][0]) + sqr(r1[1] - r2[3][1]) + sqr(r1[2] - r2[3][2]);
-            dsp[3] = sqr(sp1.x - sp2[3].x) + sqr(sp1.y - sp2[3].y); 
+            dr[3] = dist2(r1, r2[3]);
+            dsp[3] = dist2(sp1, sp2[3]);
 
-            g.addEdge(PIX(y, x), PIX(y, x + 1), SegmLinkVal(dr[0], dsp[0]));
-            g.addEdge(PIX(y, x), PIX(y + 1, x), SegmLinkVal(dr[1], dsp[1]));
-            g.addEdge(PIX(y, x), PIX(y + 1, x + 1), SegmLinkVal(dr[2], dsp[2]));
-            g.addEdge(PIX(y, x + 1), PIX(y, x + 1), SegmLinkVal(dr[3], dsp[3]));
+            g.addEdge(pix(y, x, ncols), pix(y, x + 1, ncols), SegmLinkVal(dr[0], dsp[0]));
+            g.addEdge(pix(y, x, ncols), pix(y + 1, x, ncols), SegmLinkVal(dr[1], dsp[1]));
+            g.addEdge(pix(y, x, ncols), pix(y + 1, x + 1, ncols), SegmLinkVal(dr[2], dsp[2]));
+            g.addEdge(pix(y, x + 1, ncols), pix(y, x + 1, ncols), SegmLinkVal(dr[3], dsp[3]));
         }
     }
     for (int y = 0; y < nrows - 1; ++y)
     {
         r1 = rmap.at<Vec4b>(y, ncols - 1);
         r2[0] = rmap.at<Vec4b>(y + 1, ncols - 1);
-        sp1 = spmap.at<Point_<short> >(y, ncols - 1);
-        sp2[0] = spmap.at<Point_<short> >(y + 1, ncols - 1);
-        dr[0] = sqr(r1[0] - r2[0][0]) + sqr(r1[1] - r2[0][1]) + sqr(r1[2] - r2[0][2]);
-        dsp[0] = sqr(sp1.x - sp2[0].x) + sqr(sp1.y - sp2[0].y);   
-        g.addEdge(PIX(y, ncols - 1), PIX(y + 1, ncols - 1), SegmLinkVal(dr[0], dsp[0]));
+        sp1 = spmap.at<Vec2s>(y, ncols - 1);
+        sp2[0] = spmap.at<Vec2s>(y + 1, ncols - 1);
+        dr[0] = dist2(r1, r2[0]);
+        dsp[0] = dist2(sp1, sp2[0]);
+        g.addEdge(pix(y, ncols - 1, ncols), pix(y + 1, ncols - 1, ncols), SegmLinkVal(dr[0], dsp[0]));
     }
     for (int x = 0; x < ncols - 1; ++x)
     {
         r1 = rmap.at<Vec4b>(nrows - 1, x);
         r2[0] = rmap.at<Vec4b>(nrows - 1, x + 1);
-        sp1 = spmap.at<Point_<short> >(nrows - 1, x);
-        sp2[0] = spmap.at<Point_<short> >(nrows - 1, x + 1);
-        dr[0] = sqr(r1[0] - r2[0][0]) + sqr(r1[1] - r2[0][1]) + sqr(r1[2] - r2[0][2]);
-        dsp[0] = sqr(sp1.x - sp2[0].x) + sqr(sp1.y - sp2[0].y);   
-        g.addEdge(PIX(nrows - 1, x), PIX(nrows - 1, x + 1), SegmLinkVal(dr[0], dsp[0]));
+        sp1 = spmap.at<Vec2s>(nrows - 1, x);
+        sp2[0] = spmap.at<Vec2s>(nrows - 1, x + 1);
+        dr[0] = dist2(r1, r2[0]);
+        dsp[0] = dist2(sp1, sp2[0]);
+        g.addEdge(pix(nrows - 1, x, ncols), pix(nrows - 1, x + 1, ncols), SegmLinkVal(dr[0], dsp[0]));
     }
 
-    LOG2("raginit:", clock() - start);
-    DBG(start = clock());
-
     DjSets comps(g.numv);
 
-    LOG2("djsetinit:", clock() - start);
-    DBG(start = clock());
-
     // Find adjacent components
     for (int v = 0; v < g.numv; ++v)
     {
@@ -386,15 +324,9 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
         }
     }
 
-    LOG2("findadjacent:", clock() - start);
-    DBG(start = clock());
-
     vector<SegmLink> edges;
     edges.reserve(g.numv);
 
-    LOG2("initedges:", clock() - start);
-    DBG(start = clock());
-
     // Prepare edges connecting differnet components
     for (int v = 0; v < g.numv; ++v)
     {
@@ -407,14 +339,8 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
         }
     }
 
-    LOG2("prepareforsort:", clock() - start);
-    DBG(start = clock());
-
     // Sort all graph's edges connecting differnet components (in asceding order)
-    sort(edges.begin(), edges.end(), SegmLinkCmp());
-
-    LOG2("sortedges:", clock() - start);
-    DBG(start = clock());
+    sort(edges.begin(), edges.end());
 
     // Exclude small components (starting from the nearest couple)
     for (size_t i = 0; i < edges.size(); ++i)
@@ -425,9 +351,6 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
             comps.merge(c1, c2);
     }
 
-    LOG2("excludesmall:", clock() - start);
-    DBG(start = clock());
-
     // Compute sum of the pixel's colors which are in the same segment
     Mat h_src = src;
     vector<Vec4i> sumcols(nrows * ncols, Vec4i(0, 0, 0, 0));
@@ -436,7 +359,7 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
         Vec4b* h_srcy = h_src.ptr<Vec4b>(y);
         for (int x = 0; x < ncols; ++x)
         {
-            int parent = comps.find(PIX(y, x));
+            int parent = comps.find(pix(y, x, ncols));
             Vec4b col = h_srcy[x];
             Vec4i& sumcol = sumcols[parent];
             sumcol[0] += col[0];
@@ -445,9 +368,6 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
         }
     }
 
-    LOG2("computesum:", clock() - start);
-    DBG(start = clock());
-
     // Create final image, color of each segment is the average color of its pixels
     dst.create(src.size(), src.type());
 
@@ -456,7 +376,7 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
         Vec4b* dsty = dst.ptr<Vec4b>(y);
         for (int x = 0; x < ncols; ++x)
         {
-            int parent = comps.find(PIX(y, x));
+            int parent = comps.find(pix(y, x, ncols));
             const Vec4i& sumcol = sumcols[parent];
             Vec4b& dstcol = dsty[x];
             dstcol[0] = static_cast<uchar>(sumcol[0] / comps.size[parent]);
@@ -464,11 +384,6 @@ void meanShiftSegmentation(const GpuMat& src, Mat& dst, int sp, int sr, int mins
             dstcol[2] = static_cast<uchar>(sumcol[2] / comps.size[parent]);
         }
     }
-
-    LOG2("createfinal:", clock() - start);
 }
 
-} // namespace gpu
-} // namespace cv
-
 #endif // #if !defined (HAVE_CUDA)
\ No newline at end of file