Mex wrapper for graph segmentation

Alec Jacobson

May 04, 2017

weblog/

I wrote a small mex wrapper for the graph segmentation part of the "Graph Based Image Segmentation" code. Most of the previously matlab implementations/wrappers worked on images. I want to apply this to geometry so I needed access to the graph segmentation directly. Here's the wrapper (soon to be part of gptoolbox):

// mexopts = gptoolbox_mexopts('Static',false,'Debug',true);
// mex('segment_graph.cpp',mexopts{:});
#ifdef MEX
#  include <mex.h>
#  include <igl/C_STR.h>
#  include <igl/matlab/mexErrMsgTxt.h>
#  undef assert
#  define assert( isOK ) ( (isOK) ? (void)0 : (void) ::mexErrMsgTxt(C_STR(__FILE__<<":"<<__LINE__<<": failed assertion `"<<#isOK<<"'"<<std::endl) ) )
#endif
#include "segment-graph.h"
#include <igl/matlab/mexErrMsgTxt.h>
#include <igl/matlab/parse_rhs.h>
#include <igl/unique.h>
#include <igl/matlab/prepare_lhs.h>
#include <igl/matlab/requires_arg.h>
#include <igl/matlab/validate_arg.h>
#include <igl/matlab/MexStream.h>
#include <Eigen/Sparse>
void mexFunction(
 int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
  using namespace igl::matlab;
  using namespace Eigen;
  using namespace std;
  igl::matlab::MexStream mout;        
  std::streambuf *outbuf = std::cout.rdbuf(&mout);


  mexErrMsgTxt(nrhs>0,"Too few inputs");
  mexErrMsgTxt(mxIsSparse(prhs[0]),"Matrix should be sparse");
  const mxArray * mx_data = prhs[0];
  const int m = mxGetM(mx_data);
  const int n = mxGetN(mx_data);
  mexErrMsgTxt(n == mxGetM(prhs[0]), "Matrix should be square");
  assert(mxIsSparse(mx_data));
  assert(mxGetNumberOfDimensions(mx_data) == 2);
  // TODO: It should be possible to directly load the data into the sparse
  // matrix without going through the triplets
  // Copy data immediately
  double * pr = mxGetPr(mx_data);
  mwIndex * ir = mxGetIr(mx_data);
  mwIndex * jc = mxGetJc(mx_data);
  const int num_edges = mxGetNzmax(mx_data);
  edge * edges = new edge[num_edges];
  int k = 0;
  for(int j=0; j<n;j++)
  {
    // Iterate over inside
    while(k<(int)jc[j+1])
    {
      //cout<<ir[k]<<" "<<j<<" "<<pr[k]<<endl;
      assert((int)ir[k]<m);
      assert((int)j<n);
      edges[k].a = ir[k];
      edges[k].b = j;
      edges[k].w = pr[k];
      k++;
    }
  }

  // defaults 
  int min_size = 0;
  // Threshold 
  int c = sqrt((double)n);
  {
    int i = 1;
    while(i<nrhs)
    {
      mexErrMsgTxt(mxIsChar(prhs[i]),"Parameter names should be strings");
      // Cast to char
      const char * name = mxArrayToString(prhs[i]);
      if(strcmp("Threshold",name) == 0)
      {
        requires_arg(i,nrhs,name);
        validate_arg_scalar(i,nrhs,prhs,name);
        validate_arg_double(i,nrhs,prhs,name);
        c = (double)*mxGetPr(prhs[++i]);
      }else if(strcmp("MinSize",name) == 0)
      {
        requires_arg(i,nrhs,name);
        validate_arg_scalar(i,nrhs,prhs,name);
        validate_arg_double(i,nrhs,prhs,name);
        min_size = (int)((double)*mxGetPr(prhs[++i]));
      }
      i++;
    }
  }

  universe *u = segment_graph(n, num_edges, edges, c);

  // post process small components
  for (int i = 0; i < num_edges; i++) {
    int a = u->find(edges[i].a);
    int b = u->find(edges[i].b);
    if ((a != b) && ((u->size(a) < min_size) || (u->size(b) < min_size)))
      u->join(a, b);
  }

  switch(nlhs)
  {
    case 1:
    {
      plhs[0] = mxCreateDoubleMatrix(m,1, mxREAL);
      Eigen::VectorXi C(m);
      for(int i = 0;i<m;i++)
      {
        C(i) = u->find(i);
      }
      Eigen::VectorXi uC,I,J;
      igl::unique(C,uC,I,J);
      prepare_lhs_index(J,plhs);
    }
    default: break;
  }

  delete[] edges;
  delete u;
  std::cout.rdbuf(outbuf);
}

It takes the graph as a sparse matrix and outputs the component ids:

C = segment_graph(A);