├── .gitignore ├── data ├── background.stroke ├── foreground.stroke ├── soldier_baseline.fbmask ├── soldier.png ├── background.png ├── foreground.png └── soldier_baseline.png ├── .gitmodules ├── Examples ├── CMakeLists.txt └── ImageGraphCutSegmentationExample.cpp ├── Tests ├── CMakeLists.txt └── ImageGraphCutSegmentationTest.cpp ├── README.md ├── CMakeLists.txt ├── PixelDifference.h ├── ImageGraphCut.h └── ImageGraphCut.hpp /.gitignore: -------------------------------------------------------------------------------- 1 | *.user 2 | *.user.* 3 | *~ -------------------------------------------------------------------------------- /data/background.stroke: -------------------------------------------------------------------------------- 1 | stroke 255 2 | background.png -------------------------------------------------------------------------------- /data/foreground.stroke: -------------------------------------------------------------------------------- 1 | stroke 255 2 | foreground.png -------------------------------------------------------------------------------- /data/soldier_baseline.fbmask: -------------------------------------------------------------------------------- 1 | foreground 255 2 | background 0 3 | soldier_baseline.png -------------------------------------------------------------------------------- /data/soldier.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/daviddoria/ImageGraphCutSegmentation/HEAD/data/soldier.png -------------------------------------------------------------------------------- /data/background.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/daviddoria/ImageGraphCutSegmentation/HEAD/data/background.png -------------------------------------------------------------------------------- /data/foreground.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/daviddoria/ImageGraphCutSegmentation/HEAD/data/foreground.png -------------------------------------------------------------------------------- /data/soldier_baseline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/daviddoria/ImageGraphCutSegmentation/HEAD/data/soldier_baseline.png -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "Mask"] 2 | path = Mask 3 | url = https://github.com/daviddoria/Mask.git 4 | [submodule "CMakeHelpers"] 5 | path = CMakeHelpers 6 | url = https://github.com/daviddoria/CMakeHelpers.git 7 | -------------------------------------------------------------------------------- /Examples/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | ADD_EXECUTABLE(ImageGraphCutSegmentationExample ImageGraphCutSegmentationExample.cpp) 2 | TARGET_LINK_LIBRARIES(ImageGraphCutSegmentationExample 3 | ${ImageGraphCutSegmentation_libraries}) 4 | -------------------------------------------------------------------------------- /Tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | enable_testing() 2 | 3 | ADD_EXECUTABLE(ImageGraphCutSegmentationTest ImageGraphCutSegmentationTest.cpp) 4 | TARGET_LINK_LIBRARIES(ImageGraphCutSegmentationTest 5 | ${ITK_LIBRARIES} 6 | ${ImageGraphCutSegmentation_libraries}) 7 | 8 | add_test(NAME ImageGraphCutSegmentationTest COMMAND ImageGraphCutSegmentationTest 9 | ${CMAKE_SOURCE_DIR}/data/soldier.png ${CMAKE_SOURCE_DIR}/data/foreground.stroke 10 | ${CMAKE_SOURCE_DIR}/data/background.stroke ${CMAKE_SOURCE_DIR}/data/soldier_baseline.fbmask) 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Getting the code 2 | ---------------- 3 | After you have cloned this repository, you will need to initialize the submodules: 4 | git submodule update --init --recursive 5 | 6 | If you ever update with 'git pull origin master', you must then do 'git submodule update --recursive'. 7 | 8 | This pulls in all of the dependencies including Mask (which includes ITKHelpers and then Helpers). 9 | 10 | Overview 11 | -------- 12 | This software allows the user to perform a foreground/background segmentation of an image. 13 | This implementation is based on "Graph Cuts and Efficient N-D Image Segmentation" by Yuri Boykov (IJCV 2006). 14 | 15 | Build notes 16 | ------------------ 17 | This code depends on c++0x/11 additions to the c++ language. For Linux, this means it must be built with the flag 18 | gnu++0x (or gnu++11 for gcc >= 4.7). 19 | 20 | Dependencies 21 | ------------ 22 | - ITK >= 4 23 | - Boost 1.51 24 | You can tell this project's CMake to use a local boost build with: cmake . -DBOOST_ROOT=/home/doriad/build/boost_1_51 25 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8) 2 | SET(CMAKE_INCLUDE_CURRENT_DIR ON) 3 | 4 | project(ImageGraphCutSegmentation) 5 | enable_testing() 6 | 7 | add_subdirectory(CMakeHelpers) 8 | 9 | # Compile with c++11 functionality. 10 | if(UNIX) 11 | SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -std=gnu++11") # gcc >= 4.7 12 | endif(UNIX) 13 | 14 | # ITK 15 | FIND_PACKAGE(ITK REQUIRED ITKCommon ITKIOImageBase ITKIOPNG ITKIOMeta 16 | ITKImageIntensity ITKImageFeature 17 | ITKMathematicalMorphology 18 | ITKBinaryMathematicalMorphology ITKDistanceMap 19 | ITKTestKernel) 20 | INCLUDE( ${ITK_USE_FILE} ) 21 | 22 | if( "${ITK_VERSION_MAJOR}" LESS 4 ) 23 | MESSAGE(FATAL_ERROR "You must build ImageGraphCutSegmentation with ITK >= 4.0!") 24 | endif( "${ITK_VERSION_MAJOR}" LESS 4 ) 25 | 26 | # Boost 27 | set(Boost_USE_MULTITHREADED ON) 28 | FIND_PACKAGE(Boost 1.50) 29 | INCLUDE_DIRECTORIES(${Boost_INCLUDE_DIRS}) 30 | 31 | # Submodules 32 | UseSubmodule(Mask ImageGraphCutSegmentation) 33 | 34 | # Allow this project to be detected and used as a submodule 35 | CreateSubmodule(ImageGraphCutSegmentation) 36 | 37 | # Give the compiler all of the required include directories 38 | include_directories(${ImageGraphCutSegmentation_include_dirs}) 39 | 40 | # Make the h/hpp files appear in a QtCreator project 41 | add_custom_target(ImageGraphCut SOURCES 42 | ImageGraphCut.h ImageGraphCut.hpp PixelDifference.h README.md) 43 | 44 | option(ImageGraphCutSegmentation_BuildExamples "Build ImageGraphCutSegmentation examples?") 45 | if(ImageGraphCutSegmentation_BuildExamples) 46 | add_subdirectory(Examples) 47 | endif() 48 | 49 | option(ImageGraphCutSegmentation_BuildTests "Build ImageGraphCutSegmentation tests?") 50 | if(ImageGraphCutSegmentation_BuildTests) 51 | add_subdirectory(Tests) 52 | endif() 53 | -------------------------------------------------------------------------------- /PixelDifference.h: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (C) 2012 David Doria, daviddoria@gmail.com 3 | 4 | This program is free software: you can redistribute it and/or modify 5 | it under the terms of the GNU General Public License as published by 6 | the Free Software Foundation, either version 3 of the License, or 7 | (at your option) any later version. 8 | 9 | This program is distributed in the hope that it will be useful, 10 | but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | GNU General Public License for more details. 13 | 14 | You should have received a copy of the GNU General Public License 15 | along with this program. If not, see . 16 | */ 17 | 18 | #ifndef PixelDifference_H 19 | #define PixelDifference_H 20 | 21 | #include 22 | 23 | /** Compute the difference between two RGB pixels. */ 24 | template 25 | class 26 | RGBPixelDifference 27 | { 28 | public: 29 | float Difference(const TPixel& a, const TPixel& b) 30 | { 31 | // Compute the Euclidean distance between N dimensional pixels 32 | float difference = 0; 33 | 34 | for(unsigned int i = 0; i < 3; i++) 35 | { 36 | difference += pow(a[i] - b[i], 2); 37 | } 38 | 39 | return sqrt(difference); 40 | } 41 | }; 42 | 43 | /** Compute the difference between two pixels whose first 3 components are 44 | * RGB and then there are more components afterwards. E.g. RGBD. */ 45 | template 46 | class NDPixelDifference 47 | { 48 | //float RGBWeight = 1.0f; // Needs better c++11 support than is provided by VS2010 49 | float RGBWeight; 50 | 51 | public: 52 | NDPixelDifference(const float rgbWeight) : RGBWeight(rgbWeight){} 53 | 54 | float Difference(const TPixel& a, const TPixel& b) 55 | { 56 | assert(a.GetNumberOfComponents() == b.GetNumberOfComponents()); 57 | 58 | unsigned int numberOfComponents = a.GetNumberOfComponents(); 59 | 60 | // Compute the Euclidean distance between N dimensional pixels 61 | float difference = 0; 62 | 63 | if(a.GetNumberOfComponents() > 3) 64 | { 65 | for(unsigned int i = 0; i < 3; i++) 66 | { 67 | difference += (this->RGBWeight / 3.) * pow(a[i] - b[i],2); 68 | } 69 | for(unsigned int i = 3; i < numberOfComponents; i++) 70 | { 71 | difference += (1 - this->RGBWeight) / (numberOfComponents - 3.) * pow(a[i] - b[i],2); 72 | } 73 | } 74 | else // image is RGB or less (grayscale) 75 | { 76 | for(unsigned int i = 0; i < numberOfComponents; i++) 77 | { 78 | difference += pow(a[i] - b[i],2); 79 | } 80 | } 81 | return sqrt(difference); 82 | } 83 | }; 84 | 85 | #endif 86 | -------------------------------------------------------------------------------- /Tests/ImageGraphCutSegmentationTest.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (C) 2012 David Doria, daviddoria@gmail.com 3 | 4 | This program is free software: you can redistribute it and/or modify 5 | it under the terms of the GNU General Public License as published by 6 | the Free Software Foundation, either version 3 of the License, or 7 | (at your option) any later version. 8 | 9 | This program is distributed in the hope that it will be useful, 10 | but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | GNU General Public License for more details. 13 | 14 | You should have received a copy of the GNU General Public License 15 | along with this program. If not, see . 16 | */ 17 | 18 | #include "ImageGraphCut.h" 19 | 20 | // Submodules 21 | #include "Mask/ITKHelpers/Helpers/Helpers.h" 22 | #include "Mask/ITKHelpers/ITKHelpers.h" 23 | #include "Mask/StrokeMask.h" 24 | 25 | // ITK 26 | #include "itkImage.h" 27 | #include "itkImageFileReader.h" 28 | #include "itkImageFileWriter.h" 29 | #include "itkImageRegionConstIteratorWithIndex.h" 30 | #include "itkTestingComparisonImageFilter.h" 31 | #include "itkVectorImage.h" 32 | 33 | int main(int argc, char*argv[]) 34 | { 35 | // Verify arguments 36 | if(argc != 5) 37 | { 38 | std::cerr << "Required: image.png foreground.stroke background.stroke baseline.png" << std::endl; 39 | return EXIT_FAILURE; 40 | } 41 | 42 | // Parse arguments 43 | std::string imageFilename = argv[1]; 44 | 45 | std::string foregroundStrokeFilename = argv[2]; 46 | 47 | std::string backgroundStrokeFilename = argv[3]; 48 | 49 | std::string baselineFilename = argv[4]; 50 | 51 | // Output arguments 52 | std::cout << "imageFilename: " << imageFilename << std::endl 53 | << "foregroundStrokeFilename: " << foregroundStrokeFilename << std::endl 54 | << "backgroundStrokeFilename: " << backgroundStrokeFilename << std::endl 55 | << "baselineFilename: " << baselineFilename << std::endl; 56 | 57 | // The type of the image to segment 58 | typedef itk::VectorImage ImageType; 59 | 60 | // Read the image 61 | typedef itk::ImageFileReader ReaderType; 62 | ReaderType::Pointer reader = ReaderType::New(); 63 | reader->SetFileName(imageFilename); 64 | reader->Update(); 65 | 66 | // Read the foreground and background stroke images 67 | StrokeMask::Pointer foregroundStrokeMask = 68 | StrokeMask::New(); 69 | foregroundStrokeMask->Read(foregroundStrokeFilename); 70 | 71 | StrokeMask::Pointer backgroundStrokeMask = 72 | StrokeMask::New(); 73 | backgroundStrokeMask->Read(backgroundStrokeFilename); 74 | 75 | // Perform the cut 76 | ImageGraphCut GraphCut; 77 | GraphCut.SetImage(reader->GetOutput()); 78 | GraphCut.SetNumberOfHistogramBins(20); 79 | GraphCut.SetLambda(.01); 80 | std::vector > foregroundPixels = 81 | ITKHelpers::GetPixelsWithValue(foregroundStrokeMask.GetPointer(), StrokeMaskPixelTypeEnum::STROKE); 82 | std::cout << "There are " << foregroundPixels.size() << " foreground pixels." << std::endl; 83 | 84 | std::vector > backgroundPixels = 85 | ITKHelpers::GetPixelsWithValue(backgroundStrokeMask.GetPointer(), StrokeMaskPixelTypeEnum::STROKE); 86 | std::cout << "There are " << backgroundPixels.size() << " background pixels." << std::endl; 87 | 88 | GraphCut.SetSources(foregroundPixels); 89 | GraphCut.SetSinks(backgroundPixels); 90 | GraphCut.PerformSegmentation(); 91 | 92 | // Get and write the result 93 | ForegroundBackgroundSegmentMask* result = GraphCut.GetSegmentMask(); 94 | 95 | // Read the baseline image 96 | ForegroundBackgroundSegmentMask::Pointer baselineMask = 97 | ForegroundBackgroundSegmentMask::New(); 98 | baselineMask->Read(baselineFilename); 99 | 100 | baselineMask->Write("BaselineReadWrite.png", ForegroundPixelValueWrapper(255), 101 | BackgroundPixelValueWrapper(0)); 102 | 103 | unsigned int numberOfIncorrectPixels = ITKHelpers::CountDifferentPixels(result, baselineMask.GetPointer()); 104 | if(numberOfIncorrectPixels > 0) 105 | { 106 | std::cerr << "There were " << numberOfIncorrectPixels << " incorrect pixels!" << std::endl; 107 | result->Write("Incorrect.png", ForegroundPixelValueWrapper(255), 108 | BackgroundPixelValueWrapper(0)); 109 | return EXIT_FAILURE; 110 | } 111 | 112 | return EXIT_SUCCESS; 113 | } 114 | -------------------------------------------------------------------------------- /Examples/ImageGraphCutSegmentationExample.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (C) 2012 David Doria, daviddoria@gmail.com 3 | 4 | This program is free software: you can redistribute it and/or modify 5 | it under the terms of the GNU General Public License as published by 6 | the Free Software Foundation, either version 3 of the License, or 7 | (at your option) any later version. 8 | 9 | This program is distributed in the hope that it will be useful, 10 | but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | GNU General Public License for more details. 13 | 14 | You should have received a copy of the GNU General Public License 15 | along with this program. If not, see . 16 | */ 17 | 18 | #include "ImageGraphCut.h" 19 | 20 | // Submodules 21 | #include "Mask/ITKHelpers/Helpers/Helpers.h" 22 | #include "Mask/ITKHelpers/ITKHelpers.h" 23 | #include "Mask/StrokeMask.h" 24 | 25 | #include "itkImage.h" 26 | #include "itkImageFileReader.h" 27 | #include "itkImageFileWriter.h" 28 | #include "itkImageRegionConstIteratorWithIndex.h" 29 | #include "itkVectorImage.h" 30 | 31 | /** This example segments an image and writes the segmentation mask to a file. 32 | * It can operate on image with an number of components per dimension (i.e. 33 | * RGB, RGBD, grayscale, etc.). NOTE: If your image has channels that should not 34 | * be used by the segmentation algorithm, you must remove them before calling this 35 | * example. That is, if you have an RGBA image (4th channel is an alpha channel) 36 | * but you only want to segment the RGB image, you must remove the alpha channel 37 | * or you will get strange results (the full image may be the resulting segmentation). 38 | */ 39 | int main(int argc, char*argv[]) 40 | { 41 | // Verify arguments 42 | if(argc != 5) 43 | { 44 | std::cerr << "Required: image.png foregroundMask.png backgroundMask.png output.png" << std::endl; 45 | return EXIT_FAILURE; 46 | } 47 | 48 | // Parse arguments 49 | std::string imageFilename = argv[1]; 50 | 51 | // This image should have white pixels indicating foreground pixels and be black elsewhere. 52 | std::string foregroundFilename = argv[2]; 53 | 54 | // This image should have white pixels indicating background pixels and be black elsewhere. 55 | std::string backgroundFilename = argv[3]; 56 | 57 | std::string outputFilename = argv[4]; // Foreground/background segment mask 58 | 59 | // Output arguments 60 | std::cout << "imageFilename: " << imageFilename << std::endl 61 | << "foregroundFilename: " << foregroundFilename << std::endl 62 | << "backgroundFilename: " << backgroundFilename << std::endl 63 | << "outputFilename: " << outputFilename << std::endl; 64 | 65 | // The type of the image to segment 66 | typedef itk::VectorImage ImageType; 67 | 68 | // Read the image 69 | typedef itk::ImageFileReader ReaderType; 70 | ReaderType::Pointer reader = ReaderType::New(); 71 | reader->SetFileName(imageFilename); 72 | reader->Update(); 73 | 74 | std::cout << "Read image with " << reader->GetOutput()->GetNumberOfComponentsPerPixel() 75 | << " components per pixel." << std::endl; 76 | 77 | // Read the foreground and background stroke images 78 | StrokeMask::Pointer foregroundStrokeMask = 79 | StrokeMask::New(); 80 | foregroundStrokeMask->Read(foregroundFilename); 81 | 82 | StrokeMask::Pointer backgroundStrokeMask = 83 | StrokeMask::New(); 84 | backgroundStrokeMask->Read(backgroundFilename); 85 | 86 | // Perform the cut 87 | std::cout << "Starting graphcut..." << std::endl; 88 | ImageGraphCut GraphCut; 89 | GraphCut.SetImage(reader->GetOutput()); 90 | GraphCut.SetNumberOfHistogramBins(20); 91 | GraphCut.SetLambda(.01); 92 | std::vector > foregroundPixels = 93 | ITKHelpers::GetPixelsWithValue(foregroundStrokeMask.GetPointer(), StrokeMaskPixelTypeEnum::STROKE); 94 | std::vector > backgroundPixels = 95 | ITKHelpers::GetPixelsWithValue(backgroundStrokeMask.GetPointer(), StrokeMaskPixelTypeEnum::STROKE); 96 | GraphCut.SetSources(foregroundPixels); 97 | GraphCut.SetSinks(backgroundPixels); 98 | GraphCut.PerformSegmentation(); 99 | 100 | // Get and write the result 101 | ForegroundBackgroundSegmentMask* segmentMask = GraphCut.GetSegmentMask(); 102 | 103 | segmentMask->Write("resultingMask.png", ForegroundPixelValueWrapper(0), 104 | BackgroundPixelValueWrapper(255)); 105 | 106 | ImageType::Pointer result = ImageType::New(); 107 | ITKHelpers::DeepCopy(reader->GetOutput(), result.GetPointer()); 108 | ImageType::PixelType backgroundColor(3); 109 | backgroundColor.Fill(0); 110 | segmentMask->ApplyToImage(result.GetPointer(), backgroundColor); 111 | 112 | ITKHelpers::WriteImage(result.GetPointer(), outputFilename); 113 | 114 | } 115 | -------------------------------------------------------------------------------- /ImageGraphCut.h: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (C) 2012 David Doria, daviddoria@gmail.com 3 | 4 | This program is free software: you can redistribute it and/or modify 5 | it under the terms of the GNU General Public License as published by 6 | the Free Software Foundation, either version 3 of the License, or 7 | (at your option) any later version. 8 | 9 | This program is distributed in the hope that it will be useful, 10 | but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | GNU General Public License for more details. 13 | 14 | You should have received a copy of the GNU General Public License 15 | along with this program. If not, see . 16 | */ 17 | 18 | #ifndef ImageGraphCut_H 19 | #define ImageGraphCut_H 20 | 21 | // Custom 22 | #include "PixelDifference.h" 23 | 24 | // Submodules 25 | #include "Mask/ForegroundBackgroundSegmentMask.h" 26 | 27 | // ITK 28 | #include "itkImage.h" 29 | #include "itkSampleToHistogramFilter.h" 30 | #include "itkHistogram.h" 31 | #include "itkListSample.h" 32 | 33 | // STL 34 | #include 35 | 36 | // Boost 37 | #include 38 | #include 39 | #include 40 | #include 41 | #include 42 | 43 | /** Perform graph cut based segmentation on an image. Image pixels can contain any 44 | * number of components (i.e. grayscale, RGB, RGBA, RGBD, etc.). 45 | * This is an implementation of the technique described here: 46 | * http://www.eecs.berkeley.edu/~efros/courses/AP06/Papers/boykov-iccv-01.pdf 47 | */ 48 | template > 49 | class ImageGraphCut 50 | { 51 | public: 52 | // Typedefs 53 | 54 | /** This is a special type to keep track of the graph node labels. */ 55 | typedef itk::Image NodeImageType; 56 | 57 | /** The type of the histograms. */ 58 | typedef itk::Statistics::Histogram< float, 59 | itk::Statistics::DenseFrequencyContainer2 > HistogramType; 60 | 61 | /** The type of a list of pixels/indexes. */ 62 | typedef std::vector > IndexContainer; 63 | 64 | typedef typename TImage::PixelType PixelType; 65 | typedef itk::Statistics::ListSample SampleType; 66 | typedef itk::Statistics::SampleToHistogramFilter SampleToHistogramFilterType; 67 | 68 | /** If nothing else is provided, this is the default background likelihood function. */ 69 | float InternalForegroundLikelihood(const PixelType& pixel); 70 | 71 | /** If nothing else is provided, this is the default background likelihood function. */ 72 | float InternalBackgroundLikelihood(const PixelType& pixel); 73 | 74 | ImageGraphCut() 75 | { 76 | this->ForegroundLikelihood = 77 | boost::bind( 78 | &ImageGraphCut:: 79 | InternalForegroundLikelihood, this, _1); 80 | 81 | this->BackgroundLikelihood = 82 | boost::bind( 83 | &ImageGraphCut:: 84 | InternalBackgroundLikelihood, this, _1); 85 | } 86 | 87 | ImageGraphCut(TPixelDifferenceFunctor pixelDifferenceFunctor) : 88 | PixelDifferenceFunctor(pixelDifferenceFunctor){} 89 | 90 | TPixelDifferenceFunctor PixelDifferenceFunctor; 91 | 92 | /** Provide the image to segment. */ 93 | void SetImage(TImage* const image); 94 | 95 | /** Several initializations are done here. */ 96 | void Initialize(); 97 | 98 | /** Get the image that we are segmenting. */ 99 | TImage* GetImage(); 100 | 101 | /** Create and cut the graph (The main driver function). */ 102 | void PerformSegmentation(); 103 | 104 | /** Return a list of the selected (via scribbling) pixels. */ 105 | IndexContainer GetSources(); 106 | IndexContainer GetSinks(); 107 | 108 | /** Set the selected pixels. */ 109 | void SetSources(const IndexContainer& sources); 110 | void SetSinks(const IndexContainer& sinks); 111 | 112 | /** Get the output of the segmentation. */ 113 | ForegroundBackgroundSegmentMask* GetSegmentMask(); 114 | 115 | /** Set the weight between the regional and boundary terms. */ 116 | void SetLambda(const float); 117 | 118 | /** Set the number of bins per dimension of the foreground and background histograms. */ 119 | void SetNumberOfHistogramBins(const int); 120 | 121 | void SetForegroundLikelihoodFunction(boost::function f) 122 | { 123 | this->ForegroundLikelihood = f; 124 | this->CustomLikelihood = true; 125 | } 126 | 127 | void SetBackgroundLikelihoodFunction(boost::function f) 128 | { 129 | this->BackgroundLikelihood = f; 130 | this->CustomLikelihood = true; 131 | } 132 | 133 | protected: 134 | 135 | /** A graph object for Kolmogorov*/ 136 | typedef boost::adjacency_list > GraphType; 139 | 140 | typedef boost::graph_traits::vertex_descriptor VertexDescriptor; 141 | typedef boost::graph_traits::edge_descriptor EdgeDescriptor; 142 | typedef boost::graph_traits::vertices_size_type VertexIndex; 143 | typedef boost::graph_traits::edges_size_type EdgeIndex; 144 | 145 | /** Store the list of edges and their corresponding reverse edges. */ 146 | std::vector ReverseEdges; 147 | 148 | /** Create an edge on the graph. */ 149 | unsigned int AddBidirectionalEdge(unsigned int numberOfEdges, const unsigned int source, 150 | const unsigned int target, 151 | const float weight); 152 | 153 | /** The main graph object. */ 154 | GraphType Graph; 155 | 156 | /** Maintain a list of all of the edge weights. */ 157 | std::vector EdgeWeights; 158 | 159 | /** The output segmentation */ 160 | ForegroundBackgroundSegmentMask::Pointer ResultingSegments; 161 | 162 | /** User specified foreground points */ 163 | IndexContainer Sources; 164 | 165 | /** User specified background points */ 166 | IndexContainer Sinks; 167 | 168 | /** The weighting between unary and binary terms */ 169 | float Lambda = 0.01f; 170 | 171 | /** The number of bins per dimension of the foreground and background histograms */ 172 | int NumberOfHistogramBins = 10; 173 | 174 | /** An image which keeps tracks of the mapping between pixel index and graph node id */ 175 | NodeImageType::Pointer NodeImage; 176 | 177 | /** Create the histograms from the users selections */ 178 | void CreateSamples(); 179 | 180 | /** Estimate the "camera noise" */ 181 | double ComputeNoise(); 182 | 183 | /** Create a Kolmogorov graph structure from the image and selections */ 184 | void CreateGraph(); 185 | 186 | /** Create the edges between pixels and neighboring pixels (the grid). */ 187 | void CreateNEdges(); 188 | 189 | /** Create the edges between pixels and the terminals (source and sink). */ 190 | void CreateTEdges(); 191 | 192 | /** Perform the s-t min cut */ 193 | void CutGraph(); 194 | 195 | /** The ITK data structure for storing the values that we will compute the histogram of. */ 196 | typename SampleType::Pointer ForegroundSample; 197 | typename SampleType::Pointer BackgroundSample; 198 | 199 | /** The histograms. */ 200 | const HistogramType* ForegroundHistogram = nullptr; 201 | const HistogramType* BackgroundHistogram = nullptr; 202 | 203 | /** ITK filters to create histograms. */ 204 | typename SampleToHistogramFilterType::Pointer ForegroundHistogramFilter; 205 | typename SampleToHistogramFilterType::Pointer BackgroundHistogramFilter; 206 | 207 | /** The image to be segmented */ 208 | typename TImage::Pointer Image; 209 | 210 | /** The node id of the foreground terminal. */ 211 | unsigned int SourceNodeId; 212 | 213 | /** The node id of the background terminal. */ 214 | unsigned int SinkNodeId; 215 | 216 | /** The function pointer that gets called to determine the likelihood that the pixel belongs to the foreground. */ 217 | boost::function ForegroundLikelihood; 218 | 219 | /** The function pointer that gets called to determine the likelihood that the pixel belongs to the background. */ 220 | boost::function BackgroundLikelihood; 221 | 222 | /** If a custom likelihood function is set, we don't need to compute histograms internally. */ 223 | bool CustomLikelihood = false; 224 | }; 225 | 226 | #include "ImageGraphCut.hpp" 227 | 228 | #endif 229 | -------------------------------------------------------------------------------- /ImageGraphCut.hpp: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (C) 2012 David Doria, daviddoria@gmail.com 3 | 4 | This program is free software: you can redistribute it and/or modify 5 | it under the terms of the GNU General Public License as published by 6 | the Free Software Foundation, either version 3 of the License, or 7 | (at your option) any later version. 8 | 9 | This program is distributed in the hope that it will be useful, 10 | but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | GNU General Public License for more details. 13 | 14 | You should have received a copy of the GNU General Public License 15 | along with this program. If not, see . 16 | */ 17 | 18 | #ifndef ImageGraphCut_HPP 19 | #define ImageGraphCut_HPP 20 | 21 | #include "ImageGraphCut.h" 22 | 23 | // Submodules 24 | #include "Mask/ITKHelpers/Helpers/Helpers.h" 25 | #include "Mask/ITKHelpers/ITKHelpers.h" 26 | 27 | // ITK 28 | #include "itkImageRegionIterator.h" 29 | #include "itkShapedNeighborhoodIterator.h" 30 | #include "itkMaskImageFilter.h" 31 | 32 | // STL 33 | #include 34 | #include 35 | 36 | // Boost 37 | #include 38 | 39 | template 40 | void ImageGraphCut::SetImage(TImage* const image) 41 | { 42 | this->Image = TImage::New(); 43 | ITKHelpers::DeepCopy(image, this->Image.GetPointer()); 44 | } 45 | 46 | template 47 | void ImageGraphCut::Initialize() 48 | { 49 | // Setup the output (mask) image 50 | this->ResultingSegments = ForegroundBackgroundSegmentMask::New(); 51 | this->ResultingSegments->SetRegions(this->Image->GetLargestPossibleRegion()); 52 | this->ResultingSegments->Allocate(); 53 | 54 | // Setup the image to store the node ids 55 | this->NodeImage = NodeImageType::New(); 56 | this->NodeImage->SetRegions(this->Image->GetLargestPossibleRegion()); 57 | this->NodeImage->Allocate(); 58 | 59 | // Initializations 60 | this->ForegroundSample = SampleType::New(); 61 | this->BackgroundSample = SampleType::New(); 62 | 63 | this->ForegroundHistogramFilter = SampleToHistogramFilterType::New(); 64 | this->BackgroundHistogramFilter = SampleToHistogramFilterType::New(); 65 | 66 | // Blank the NodeImage 67 | ITKHelpers::SetImageToConstant(this->NodeImage.GetPointer(), 0); 68 | 69 | // Blank the output image 70 | ITKHelpers::SetImageToConstant(this->ResultingSegments.GetPointer(), 71 | ForegroundBackgroundSegmentMaskPixelTypeEnum::BACKGROUND); 72 | 73 | } 74 | 75 | template 76 | void ImageGraphCut::CutGraph() 77 | { 78 | // Compute mininum cut 79 | std::cout << "CutGraph()..." << std::endl; 80 | boost::graph_traits::vertex_descriptor s = vertex(this->SourceNodeId, this->Graph); 81 | boost::graph_traits::vertex_descriptor t = vertex(this->SinkNodeId, this->Graph); 82 | 83 | std::vector groups(num_vertices(this->Graph)); 84 | std::vector residual_capacity(num_edges(this->Graph)); //this needs to be initialized to 0 85 | 86 | boykov_kolmogorov_max_flow(this->Graph, 87 | boost::make_iterator_property_map(&EdgeWeights[0], get(boost::edge_index, this->Graph)), 88 | boost::make_iterator_property_map(&residual_capacity[0], get(boost::edge_index, this->Graph)), 89 | boost::make_iterator_property_map(&ReverseEdges[0], get(boost::edge_index, this->Graph)), 90 | boost::make_iterator_property_map(&groups[0], get(boost::vertex_index, this->Graph)), 91 | get(boost::vertex_index, this->Graph), 92 | s, 93 | t); 94 | 95 | std::cout << "Finished max_flow()." << std::endl; 96 | 97 | // Iterate over the node image, querying the graph object for the association of each pixel and storing them as the output mask 98 | itk::ImageRegionConstIterator 99 | nodeImageIterator(this->NodeImage, this->NodeImage->GetLargestPossibleRegion()); 100 | nodeImageIterator.GoToBegin(); 101 | 102 | // From the documentation: 103 | // http://www.boost.org/doc/libs/1_55_0/libs/graph/doc/boykov_kolmogorov_max_flow.html 104 | // If the color of a vertex after running the algorithm is black the vertex belongs to 105 | // the source tree else it belongs to the sink-tree (used for minimum cuts). 106 | while(!nodeImageIterator.IsAtEnd()) 107 | { 108 | if(groups[nodeImageIterator.Get()] == groups[this->SourceNodeId]) 109 | { 110 | this->ResultingSegments->SetPixel(nodeImageIterator.GetIndex(), 111 | ForegroundBackgroundSegmentMaskPixelTypeEnum::FOREGROUND); 112 | } 113 | else 114 | { 115 | this->ResultingSegments->SetPixel(nodeImageIterator.GetIndex(), 116 | ForegroundBackgroundSegmentMaskPixelTypeEnum::BACKGROUND); 117 | } 118 | ++nodeImageIterator; 119 | } 120 | 121 | std::cout << "Finished CutGraph()." << std::endl; 122 | } 123 | 124 | // This function assumes that the ReverseEdges and EdgeWeights members are already large enough to accept the 125 | // new data (otherwise we would have to push_back/resize millions of times). 126 | template 127 | unsigned int ImageGraphCut:: 128 | AddBidirectionalEdge(unsigned int numberOfEdges, const unsigned int source, const unsigned int target, const float weight) 129 | { 130 | // skip "edge already exists check" if the node ids are too high (because it segfaults) 131 | if(source < num_vertices(this->Graph)+1 && target < num_vertices(this->Graph)+1) 132 | { 133 | // If the edge already exists, do nothing 134 | //auto edgeCheck = boost::edge(source, target, this->Graph); 135 | // std::cout << "Checking edge..." << std::endl; 136 | std::pair edgeCheck = boost::edge(source, target, this->Graph); 137 | // std::cout << "Edge check complete." << std::endl; 138 | if(edgeCheck.second == true) 139 | { 140 | return numberOfEdges; 141 | } 142 | } 143 | 144 | // Add edges between grid vertices. We have to create the edge and the reverse edge, 145 | // then add the reverseEdge as the corresponding reverse edge to 'edge', and then add 'edge' 146 | // as the corresponding reverse edge to 'reverseEdge' 147 | // int nextEdgeId = num_edges(this->Graph); // Calling this every time is VERY slow 148 | int nextEdgeId = numberOfEdges; 149 | 150 | EdgeDescriptor edge; 151 | bool inserted; 152 | 153 | boost::tie(edge,inserted) = add_edge(source, target, nextEdgeId, this->Graph); 154 | if(!inserted) 155 | { 156 | std::cerr << "Not inserted!" << std::endl; 157 | return numberOfEdges; 158 | } 159 | 160 | EdgeDescriptor reverseEdge = add_edge(target, source, nextEdgeId + 1, this->Graph).first; 161 | 162 | this->ReverseEdges[nextEdgeId] = reverseEdge; 163 | this->ReverseEdges[nextEdgeId + 1] = edge; 164 | this->EdgeWeights[nextEdgeId] = weight; 165 | this->EdgeWeights[nextEdgeId + 1] = weight; 166 | 167 | return numberOfEdges + 2; 168 | } 169 | 170 | template 171 | void ImageGraphCut::PerformSegmentation() 172 | { 173 | // This function performs some initializations and then creates and cuts the graph 174 | 175 | this->Initialize(); 176 | 177 | this->CreateGraph(); 178 | this->CutGraph(); 179 | } 180 | 181 | template 182 | void ImageGraphCut::CreateSamples() 183 | { 184 | // This function creates ITK samples from the scribbled pixels and then computes the foreground and background histograms 185 | std::cout << "CreateSamples()" << std::endl; 186 | 187 | // Ensure at least one pixel has been specified for both the foreground and background 188 | std::cout << "Currently there are " << this->Sources.size() << " sources and " 189 | << this->Sinks.size() << " sinks." << std::endl; 190 | if((this->Sources.size() <= 0) || (this->Sinks.size() <= 0)) 191 | { 192 | std::cerr << "At least one source (foreground) pixel and one sink (background) " 193 | "pixel must be specified!" << std::endl; 194 | return; 195 | } 196 | 197 | unsigned int numberOfComponentsPerPixel = 198 | this->Image->GetNumberOfComponentsPerPixel(); 199 | 200 | // We want the histogram bins to take values from 0 to 255 in all dimensions 201 | HistogramType::MeasurementVectorType binMinimum(numberOfComponentsPerPixel); 202 | HistogramType::MeasurementVectorType binMaximum(numberOfComponentsPerPixel); 203 | for(unsigned int i = 0; i < numberOfComponentsPerPixel; i++) 204 | { 205 | // If one channel (often, the alpha channel) is all 255, for example, then however 206 | // the ITK histogram constructs the bins with range (0,255) does not include the 255 pixels, 207 | // and therefore none of the pixels are included in the histogram at all! 208 | // Using slightly less than 0 (-.5) and slightly more than 255 (255.5) as the bin limits 209 | // fixes this problem 210 | // binMinimum[i] = 0; 211 | // binMaximum[i] = 255; 212 | binMinimum[i] = -0.5f; 213 | binMaximum[i] = 255.5f; 214 | } 215 | 216 | // Setup the histogram size 217 | // std::cout << "Image components per pixel: " 218 | // << numberOfComponentsPerPixel << std::endl; 219 | typename SampleToHistogramFilterType::HistogramSizeType 220 | histogramSize(numberOfComponentsPerPixel); 221 | histogramSize.Fill(this->NumberOfHistogramBins); 222 | 223 | // Create foreground samples and histogram 224 | this->ForegroundSample->Clear(); 225 | this->ForegroundSample->SetMeasurementVectorSize(numberOfComponentsPerPixel); 226 | //std::cout << "Measurement vector size: " << this->ForegroundSample->GetMeasurementVectorSize() << std::endl; 227 | //std::cout << "Pixel size: " << this->Image->GetPixel(this->Sources[0]).GetNumberOfElements() << std::endl; 228 | 229 | for(unsigned int i = 0; i < this->Sources.size(); i++) 230 | { 231 | this->ForegroundSample->PushBack(this->Image->GetPixel(this->Sources[i])); 232 | } 233 | 234 | this->ForegroundHistogramFilter->SetHistogramSize(histogramSize); 235 | this->ForegroundHistogramFilter->SetHistogramBinMinimum(binMinimum); 236 | this->ForegroundHistogramFilter->SetHistogramBinMaximum(binMaximum); 237 | this->ForegroundHistogramFilter->SetAutoMinimumMaximum(false); 238 | this->ForegroundHistogramFilter->SetInput(this->ForegroundSample); 239 | this->ForegroundHistogramFilter->Modified(); 240 | this->ForegroundHistogramFilter->Update(); 241 | 242 | this->ForegroundHistogram = this->ForegroundHistogramFilter->GetOutput(); 243 | 244 | // Create background samples and histogram 245 | this->BackgroundSample->Clear(); 246 | this->BackgroundSample->SetMeasurementVectorSize(numberOfComponentsPerPixel); 247 | for(unsigned int i = 0; i < this->Sinks.size(); i++) 248 | { 249 | this->BackgroundSample->PushBack(this->Image->GetPixel(this->Sinks[i])); 250 | } 251 | 252 | this->BackgroundHistogramFilter->SetHistogramSize(histogramSize); 253 | this->BackgroundHistogramFilter->SetHistogramBinMinimum(binMinimum); 254 | this->BackgroundHistogramFilter->SetHistogramBinMaximum(binMaximum); 255 | this->BackgroundHistogramFilter->SetAutoMinimumMaximum(false); 256 | this->BackgroundHistogramFilter->SetInput(this->BackgroundSample); 257 | this->BackgroundHistogramFilter->Modified(); 258 | this->BackgroundHistogramFilter->Update(); 259 | 260 | this->BackgroundHistogram = BackgroundHistogramFilter->GetOutput(); 261 | 262 | std::cout << "Finished CreateSamples()" << std::endl; 263 | } 264 | 265 | template 266 | void ImageGraphCut::CreateNEdges() 267 | { 268 | std::cout << "CreateNEdges()" << std::endl; 269 | // Create n-edges and set n-edge weights (links between image nodes) 270 | 271 | itk::Size<2> imageSize = this->Image->GetLargestPossibleRegion().GetSize(); 272 | 273 | unsigned int expectedNumberOfNEdges = 2*( 274 | imageSize[0] * (imageSize[1] - 1) + // vertical edges 275 | (imageSize[0]-1) * imageSize[1] // horizontal edges 276 | ); 277 | std::cout << "Resizing for " << expectedNumberOfNEdges << " N-edges." << std::endl; 278 | this->EdgeWeights.resize(num_edges(this->Graph) + expectedNumberOfNEdges); 279 | this->ReverseEdges.resize(num_edges(this->Graph) + expectedNumberOfNEdges); 280 | 281 | // We are only using a 4-connected structure, 282 | // so the kernel (iteration neighborhood) must only be 283 | // 3x3 (specified by a radius of 1) 284 | itk::Size<2> radius; 285 | radius.Fill(1); 286 | 287 | typedef itk::ShapedNeighborhoodIterator IteratorType; 288 | 289 | // Traverse the image adding an edge between the current pixel 290 | // and the pixel below it and the current pixel and the pixel to the right of it. 291 | // This prevents duplicate edges (i.e. we cannot add an edge to 292 | // all 4-connected neighbors of every pixel or almost every edge would be duplicated. 293 | std::vector neighbors; 294 | typename IteratorType::OffsetType bottom = {{0,1}}; 295 | neighbors.push_back(bottom); 296 | typename IteratorType::OffsetType right = {{1,0}}; 297 | neighbors.push_back(right); 298 | 299 | typename IteratorType::OffsetType center = {{0,0}}; 300 | 301 | IteratorType iterator(radius, this->Image, this->Image->GetLargestPossibleRegion()); 302 | iterator.ClearActiveList(); 303 | iterator.ActivateOffset(bottom); 304 | iterator.ActivateOffset(right); 305 | iterator.ActivateOffset(center); 306 | 307 | // Estimate the "camera noise" 308 | double sigma = this->ComputeNoise(); 309 | 310 | unsigned int currentNumberOfEdges = num_edges(this->Graph); 311 | for(iterator.GoToBegin(); !iterator.IsAtEnd(); ++iterator) 312 | { 313 | PixelType centerPixel = iterator.GetPixel(center); 314 | 315 | for(unsigned int i = 0; i < neighbors.size(); i++) 316 | { 317 | bool valid; 318 | iterator.GetPixel(neighbors[i], valid); 319 | 320 | // If the current neighbor is outside the image, skip it 321 | if(!valid) 322 | { 323 | continue; 324 | } 325 | 326 | PixelType neighborPixel = iterator.GetPixel(neighbors[i]); 327 | 328 | // Compute the Euclidean distance between the pixel intensities 329 | float pixelDifference = PixelDifferenceFunctor.Difference(centerPixel, neighborPixel); 330 | 331 | // Compute the edge weight 332 | float weight = exp(-pow(pixelDifference,2)/(2.0*sigma*sigma)); 333 | assert(weight >= 0); 334 | 335 | // Add the edge to the graph 336 | unsigned int node1 = this->NodeImage->GetPixel(iterator.GetIndex(center)); 337 | unsigned int node2 = this->NodeImage->GetPixel(iterator.GetIndex(neighbors[i])); 338 | 339 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, node1, node2, weight); 340 | } 341 | } 342 | 343 | std::cout << "Finished CreateNEdges()" << std::endl; 344 | } 345 | 346 | template 347 | float ImageGraphCut::InternalForegroundLikelihood(const PixelType& pixel) 348 | { 349 | HistogramType::MeasurementVectorType measurementVector(pixel.Size()); 350 | for(unsigned int i = 0; i < pixel.Size(); i++) 351 | { 352 | measurementVector[i] = pixel[i]; 353 | } 354 | 355 | HistogramType::IndexType foregroundIndex; 356 | this->ForegroundHistogram->GetIndex(measurementVector, foregroundIndex); 357 | float sourceHistogramValue = 358 | this->ForegroundHistogram->GetFrequency(foregroundIndex); 359 | 360 | sourceHistogramValue /= this->ForegroundHistogram->GetTotalFrequency(); 361 | 362 | return sourceHistogramValue; 363 | } 364 | 365 | template 366 | float ImageGraphCut::InternalBackgroundLikelihood(const PixelType& pixel) 367 | { 368 | HistogramType::MeasurementVectorType measurementVector(pixel.Size()); 369 | for(unsigned int i = 0; i < pixel.Size(); i++) 370 | { 371 | measurementVector[i] = pixel[i]; 372 | } 373 | 374 | HistogramType::IndexType backgroundIndex; 375 | this->BackgroundHistogram->GetIndex(measurementVector, backgroundIndex); 376 | float sinkHistogramValue = 377 | this->BackgroundHistogram->GetFrequency(backgroundIndex); 378 | 379 | sinkHistogramValue /= this->BackgroundHistogram->GetTotalFrequency(); 380 | 381 | return sinkHistogramValue; 382 | } 383 | 384 | template 385 | void ImageGraphCut::CreateTEdges() 386 | { 387 | // Setup links from pixel nodes to terminal nodes 388 | std::cout << "CreateTEdges()" << std::endl; 389 | 390 | itk::Size<2> imageSize = this->Image->GetLargestPossibleRegion().GetSize(); 391 | 392 | unsigned int expectedNumberOfTEdges = 2*2*(imageSize[0] * imageSize[1]); 393 | 394 | this->EdgeWeights.resize(num_edges(this->Graph) + expectedNumberOfTEdges); 395 | this->ReverseEdges.resize(num_edges(this->Graph) + expectedNumberOfTEdges); 396 | 397 | // Add t-edges and set t-edge weights (links from image nodes to virtual background and virtual foreground node) 398 | 399 | // Compute the histograms of the selected foreground and background pixels 400 | if(!this->CustomLikelihood) 401 | { 402 | CreateSamples(); 403 | } 404 | 405 | itk::ImageRegionIteratorWithIndex 406 | imageIterator(this->Image, 407 | this->Image->GetLargestPossibleRegion()); 408 | itk::ImageRegionIterator 409 | nodeIterator(this->NodeImage, 410 | this->NodeImage->GetLargestPossibleRegion()); 411 | imageIterator.GoToBegin(); 412 | nodeIterator.GoToBegin(); 413 | 414 | // Since the t-weight function takes the log of the histogram value, 415 | // we must handle bins with frequency = 0 specially (because log(0) = -inf) 416 | // For empty histogram bins we use tinyValue instead of 0. 417 | float tinyValue = 1e-10; 418 | 419 | unsigned int currentNumberOfEdges = num_edges(this->Graph); 420 | 421 | while(!imageIterator.IsAtEnd()) 422 | { 423 | itk::Index<2> currentIndex = imageIterator.GetIndex(); 424 | 425 | // Skip the computation and edge creation if the current pixel already has a fixed assignment 426 | if(Helpers::Contains(this->Sinks, currentIndex) || Helpers::Contains(this->Sources, currentIndex)) 427 | { 428 | ++imageIterator; 429 | ++nodeIterator; 430 | continue; 431 | } 432 | PixelType pixel = imageIterator.Get(); 433 | //std::cout << "Pixels have size: " << pixel.Size() << std::endl; 434 | 435 | HistogramType::MeasurementVectorType measurementVector(pixel.Size()); 436 | for(unsigned int i = 0; i < pixel.Size(); i++) 437 | { 438 | measurementVector[i] = pixel[i]; 439 | } 440 | 441 | float sourceLikelihood = ForegroundLikelihood(pixel); 442 | float sinkLikelihood = BackgroundLikelihood(pixel); 443 | 444 | if(sourceLikelihood <= 0) 445 | { 446 | sourceLikelihood = tinyValue; 447 | } 448 | 449 | if(sinkLikelihood <= 0) 450 | { 451 | sinkLikelihood = tinyValue; 452 | } 453 | 454 | // Add the edge to the graph and set its weight 455 | // log() is the natural log 456 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, nodeIterator.Get(), 457 | this->SinkNodeId, -this->Lambda*log(sourceLikelihood)); 458 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, nodeIterator.Get(), 459 | this->SourceNodeId, -this->Lambda*log(sinkLikelihood)); 460 | 461 | ++imageIterator; 462 | ++nodeIterator; 463 | } 464 | 465 | // Set very high source weights for the pixels that were 466 | // selected as foreground by the user 467 | for(unsigned int i = 0; i < this->Sources.size(); i++) 468 | { 469 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, this->NodeImage->GetPixel(this->Sources[i]), 470 | this->SourceNodeId, std::numeric_limits::max()); 471 | 472 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, this->NodeImage->GetPixel(this->Sources[i]), 473 | this->SinkNodeId, 0); 474 | } 475 | 476 | // Set very high sink weights for the pixels that 477 | // were selected as background by the user 478 | for(unsigned int i = 0; i < this->Sinks.size(); i++) 479 | { 480 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, this->NodeImage->GetPixel(this->Sinks[i]), 481 | this->SourceNodeId, 0); 482 | 483 | currentNumberOfEdges = AddBidirectionalEdge(currentNumberOfEdges, this->NodeImage->GetPixel(this->Sinks[i]), 484 | this->SinkNodeId, std::numeric_limits::max()); 485 | } 486 | 487 | std::cout << "Finished CreateTEdges()" << std::endl; 488 | } 489 | 490 | template 491 | void ImageGraphCut::CreateGraph() 492 | { 493 | std::cout << "CreateGraph()" << std::endl; 494 | 495 | // Add all of the nodes to the graph and store their IDs in a "node image" 496 | itk::ImageRegionIterator nodeImageIterator(this->NodeImage, 497 | this->NodeImage->GetLargestPossibleRegion()); 498 | nodeImageIterator.GoToBegin(); 499 | 500 | unsigned int nodeId = 0; 501 | while(!nodeImageIterator.IsAtEnd()) 502 | { 503 | nodeImageIterator.Set(nodeId); 504 | nodeId++; 505 | ++nodeImageIterator; 506 | } 507 | 508 | // Set the sink and source ids to be the two numbers immediately following the number of vertices in the grid 509 | this->SinkNodeId = nodeId; 510 | nodeId++; 511 | this->SourceNodeId = nodeId; 512 | 513 | CreateNEdges(); 514 | CreateTEdges(); 515 | 516 | { 517 | itk::Size<2> imageSize = this->NodeImage->GetLargestPossibleRegion().GetSize(); 518 | 519 | std::cout << "Number of edges " << num_edges(this->Graph) << std::endl; 520 | int expectedEdges = imageSize[0]*imageSize[1] * 2 * 2 + // one '2' is because there is an edge to both the source and sink from each pixel, and the other '2' is because they are double edges (bidirectional) 521 | 2*(imageSize[0]-1)*imageSize[1] + 2*imageSize[0]*(imageSize[1]-1); // both '2's are for the double edges (this is the number of horizontal edges + the number of vertical edges) 522 | std::cout << "(Should be " << expectedEdges << " edges.)" << std::endl; 523 | } 524 | } 525 | 526 | template 527 | double ImageGraphCut::ComputeNoise() 528 | { 529 | // Compute an estimate of the "camera noise". This is used in the N-weight function. 530 | 531 | // Since we use a 4-connected neighborhood, the kernel must be 3x3 (a rectangular radius of 1 creates a kernel side length of 3) 532 | itk::Size<2> radius; 533 | radius[0] = 1; 534 | radius[1] = 1; 535 | 536 | typedef itk::ShapedNeighborhoodIterator IteratorType; 537 | 538 | std::vector neighbors; 539 | typename IteratorType::OffsetType bottom = {{0,1}}; 540 | neighbors.push_back(bottom); 541 | typename IteratorType::OffsetType right = {{1,0}}; 542 | neighbors.push_back(right); 543 | 544 | typename IteratorType::OffsetType center = {{0,0}}; 545 | 546 | IteratorType iterator(radius, this->Image, this->Image->GetLargestPossibleRegion()); 547 | iterator.ClearActiveList(); 548 | iterator.ActivateOffset(bottom); 549 | iterator.ActivateOffset(right); 550 | iterator.ActivateOffset(center); 551 | 552 | double sigma = 0.0; 553 | int numberOfEdges = 0; 554 | 555 | // Traverse the image collecting the differences between neighboring pixel intensities 556 | for(iterator.GoToBegin(); !iterator.IsAtEnd(); ++iterator) 557 | { 558 | PixelType centerPixel = iterator.GetPixel(center); 559 | 560 | for(unsigned int i = 0; i < neighbors.size(); i++) 561 | { 562 | bool valid; 563 | iterator.GetPixel(neighbors[i], valid); 564 | if(!valid) 565 | { 566 | continue; 567 | } 568 | 569 | PixelType neighborPixel = iterator.GetPixel(neighbors[i]); 570 | 571 | float colorDifference = PixelDifferenceFunctor.Difference(centerPixel, neighborPixel); 572 | sigma += colorDifference; 573 | numberOfEdges++; 574 | } 575 | } 576 | 577 | // Normalize 578 | sigma /= static_cast(numberOfEdges); 579 | 580 | return sigma; 581 | } 582 | 583 | template 584 | typename ImageGraphCut::IndexContainer ImageGraphCut::GetSources() 585 | { 586 | return this->Sources; 587 | } 588 | 589 | template 590 | void ImageGraphCut::SetLambda(const float lambda) 591 | { 592 | this->Lambda = lambda; 593 | } 594 | 595 | template 596 | void ImageGraphCut::SetNumberOfHistogramBins(int bins) 597 | { 598 | this->NumberOfHistogramBins = bins; 599 | } 600 | 601 | template 602 | ForegroundBackgroundSegmentMask* ImageGraphCut::GetSegmentMask() 603 | { 604 | return this->ResultingSegments; 605 | } 606 | 607 | template 608 | typename ImageGraphCut::IndexContainer ImageGraphCut::GetSinks() 609 | { 610 | return this->Sinks; 611 | } 612 | 613 | template 614 | void ImageGraphCut::SetSources(const IndexContainer& sources) 615 | { 616 | this->Sources = sources; 617 | } 618 | 619 | template 620 | void ImageGraphCut::SetSinks(const IndexContainer& sinks) 621 | { 622 | this->Sinks = sinks; 623 | } 624 | 625 | template 626 | TImage* ImageGraphCut::GetImage() 627 | { 628 | return this->Image; 629 | } 630 | 631 | #endif 632 | --------------------------------------------------------------------------------