greedy-plane-segmentation-ex2.cpp

#include <iostream>
#include <vector>
#include <set>
#include <map>
#include <queue>
#include <boost/program_options/options_description.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>

#include <QtGui/qapplication.h>
#include "DGtal/base/Common.h"
#include "DGtal/io/readers/VolReader.h"
#include "DGtal/io/viewers/Viewer3D.h"
#include "DGtal/io/Display3D.h"

#include "DGtal/io/DrawWithDisplay3DModifier.h"
#include "DGtal/images/ImageSelector.h"
#include "DGtal/images/imagesSetsUtils/SetFromImage.h"
#include "DGtal/topology/DigitalSurface.h"
#include "DGtal/topology/DigitalSetBoundary.h"
#include "DGtal/topology/BreadthFirstVisitor.h"
#include "DGtal/geometry/surfaces/COBANaivePlane.h"
#include "DGtal/helpers/StdDefs.h"
#include "ConfigExamples.h"


using namespace std;
using namespace DGtal;
namespace po = boost::program_options;

using namespace Z3i;
typedef DGtal::int64_t InternalInteger;
typedef COBANaivePlane<Z3,InternalInteger> NaivePlaneComputer;
// We choose the DigitalSetBoundary surface container in order to
// segment connected or unconnected surfaces.
typedef DigitalSetBoundary<KSpace,DigitalSet> MyDigitalSurfaceContainer;
typedef DigitalSurface<MyDigitalSurfaceContainer> MyDigitalSurface;
typedef MyDigitalSurface::ConstIterator ConstIterator;
typedef MyDigitalSurface::Vertex Vertex;
typedef MyDigitalSurface::SurfelSet SurfelSet;
typedef SurfelSet::iterator SurfelSetIterator;
typedef BreadthFirstVisitor<MyDigitalSurface> Visitor;
struct SegmentedPlane {
  NaivePlaneComputer plane;
  Color color;
};
struct VertexSize {
  Vertex v;
  unsigned int size;
  inline VertexSize( const Vertex & aV, unsigned int aSize )
    : v( aV ), size( aSize ) 
  {}
};

inline
bool operator<( const VertexSize & vs1, const VertexSize & vs2 )
{
  return vs1.size < vs2.size;
} 


int main( int argc, char** argv )
{
  // parse command line ----------------------------------------------
  po::options_description general_opt("Allowed options are: ");
  general_opt.add_options()
    ("help,h", "display this message")
    ("input-file,i", po::value<std::string>()->default_value( examplesPath + "samples/Al.100.vol" ), "the volume file (.vol)" )
    ("threshold,t",  po::value<unsigned int>()->default_value(1), "the value that defines the isosurface in the image (an integer between 0 and 255)." )
    ("width-num,w",  po::value<unsigned int>()->default_value(1), "the numerator of the rational width (a non-null integer)." )
    ("width-den,d",  po::value<unsigned int>()->default_value(1), "the denominator of the rational width (a non-null integer)." );
  
  bool parseOK = true;
  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, general_opt), vm);  
  } catch ( const std::exception & ex ) {
    parseOK = false;
    trace.info() << "Error checking program options: "<< ex.what()<< endl;
  }
  po::notify(vm);
  if ( ! parseOK || vm.count("help") || ( argc <= 1 ) )
    {
      std::cout << "Usage: " << argv[0]
                << " [-i <fileName.vol>] [-t <threshold>] [-w <num>] [-d <den>]" << std::endl
                << "Segments the surface at given threshold within given volume into digital planes of rational width num/den." << std::endl
                << general_opt << std::endl;
      return 0;
    }
  string inputFilename = vm["input-file"].as<std::string>();
  unsigned int threshold = vm["threshold"].as<unsigned int>();
  unsigned int widthNum = vm["width-num"].as<unsigned int>();
  unsigned int widthDen = vm["width-den"].as<unsigned int>();

  QApplication application(argc,argv);
  typedef ImageSelector < Domain, int>::Type Image;
  Image image = VolReader<Image>::importVol(inputFilename);
  DigitalSet set3d (image.domain());
  SetFromImage<DigitalSet>::append<Image>(set3d, image, threshold,255);

  trace.beginBlock( "Set up digital surface." );
  // We initializes the cellular grid space used for defining the
  // digital surface.
  KSpace ks;
  bool ok = ks.init( set3d.domain().lowerBound(),
                     set3d.domain().upperBound(), true );
  if ( ! ok ) std::cerr << "[KSpace.init] Failed." << std::endl;
  SurfelAdjacency<KSpace::dimension> surfAdj( true ); // interior in all directions.
  MyDigitalSurfaceContainer* ptrSurfContainer = 
    new MyDigitalSurfaceContainer( ks, set3d, surfAdj );
  MyDigitalSurface digSurf( ptrSurfContainer ); // acquired
  trace.endBlock();

  Point p;
  Dimension axis;
  unsigned int j = 0;
  unsigned int nb = digSurf.size();

  // First pass to find biggest planes.
  trace.beginBlock( "1) Segmentation first pass. Computes all planes so as to sort vertices by the plane size." );
  std::map<Vertex,unsigned int> v2size;
  NaivePlaneComputer planeComputer;
  std::priority_queue<VertexSize> Q;
  for ( ConstIterator it = digSurf.begin(), itE= digSurf.end(); it != itE; ++it )
    {
      if ( ( (++j) % 50 == 0 ) || ( j == nb ) ) trace.progressBar( j, nb );
      Vertex v = *it;
      axis = ks.sOrthDir( v );
      planeComputer.init( axis, 500, widthNum, widthDen );
      // The visitor takes care of all the breadth-first traversal.
      Visitor visitor( digSurf, v );
      while ( ! visitor.finished() )
        {
          Visitor::Node node = visitor.current();
          v = node.first;
          axis = ks.sOrthDir( v );
          p = ks.sCoords( ks.sDirectIncident( v, axis ) );
          bool isExtended = planeComputer.extend( p );
          if ( isExtended ) 
            // surfel is in plane.
            visitor.expand();
          else // surfel is not in plane and should not be used in the visit.
            visitor.ignore();
        }
      Q.push( VertexSize( v, planeComputer.size() ) );
    }
  trace.endBlock();

  trace.beginBlock( "2) Segmentation second pass. Visits vertices from the one with biggest plane to the one with smallest plane." );
  std::set<Vertex> processedVertices;
  std::vector<SegmentedPlane*> segmentedPlanes;
  std::map<Vertex,SegmentedPlane*> v2plane;
  j = 0;
  while ( ! Q.empty() )
    {
      if ( ( (++j) % 50 == 0 ) || ( j == nb ) ) trace.progressBar( j, nb );
      Vertex v = Q.top().v;
      Q.pop();
      if ( processedVertices.find( v ) != processedVertices.end() ) // already in set
        continue; // process to next vertex

      SegmentedPlane* ptrSegment = new SegmentedPlane;
      segmentedPlanes.push_back( ptrSegment ); // to delete them afterwards.
      axis = ks.sOrthDir( v );
      ptrSegment->plane.init( axis, 500, widthNum, widthDen );
      // The visitor takes care of all the breadth-first traversal.
      Visitor visitor( digSurf, v );
      while ( ! visitor.finished() )
        {
          Visitor::Node node = visitor.current();
          v = node.first;
          if ( processedVertices.find( v ) == processedVertices.end() )
            { // Vertex is not in processedVertices
              axis = ks.sOrthDir( v );
              p = ks.sCoords( ks.sDirectIncident( v, axis ) );
              bool isExtended = ptrSegment->plane.extend( p );
              if ( isExtended ) 
                { // surfel is in plane.
                  processedVertices.insert( v );
                  v2plane[ v ] = ptrSegment;
                  visitor.expand();
                }
              else // surfel is not in plane and should not be used in the visit.
                visitor.ignore();
            }
          else // surfel is already in some plane.
            visitor.ignore();
        }
      // Assign random color for each plane.
      ptrSegment->color = Color( random() % 256, random() % 256, random() % 256, 255 );
    }
  trace.endBlock();

  Viewer3D viewer;
  viewer.show(); 
  for ( std::map<Vertex,SegmentedPlane*>::const_iterator 
          it = v2plane.begin(), itE = v2plane.end();
        it != itE; ++it )
    {
      viewer << CustomColors3D( it->second->color, it->second->color );
      viewer << ks.unsigns( it->first );
    }
  viewer << Display3D::updateDisplay;

  for ( std::vector<SegmentedPlane*>::iterator 
          it = segmentedPlanes.begin(), itE = segmentedPlanes.end(); 
        it != itE; ++it )
    delete *it;
  segmentedPlanes.clear();
  v2plane.clear();

  return application.exec();
}