1 /*=========================================================================
2  *
3  *  Copyright Insight Software Consortium
4  *
5  *  Licensed under the Apache License, Version 2.0 (the "License");
6  *  you may not use this file except in compliance with the License.
7  *  You may obtain a copy of the License at
8  *
9  *         http://www.apache.org/licenses/LICENSE-2.0.txt
10  *
11  *  Unless required by applicable law or agreed to in writing, software
12  *  distributed under the License is distributed on an "AS IS" BASIS,
13  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *  See the License for the specific language governing permissions and
15  *  limitations under the License.
16  *
17  *=========================================================================*/
18 
19 //
20 // This is an example of the itk::BayesianClassifierInitializationImageFilter.
21 // The example's goal is to serve as an initializer for the
22 // BayesianClassifier.cxx example also found in this directory.
23 //
24 // This example takes an input image (to be classified) and generates membership
25 // images. The membership images determine the degree to which each pixel
26 // belongs to a class.
27 //
28 // The membership image generated by the filter is an
29 // itk::VectorImage, (with pixels organized as follows: For a 2D image,
30 // its essentially a 3D array on file with DataType[y][x][c] where c is the
31 // number of classes and DataType is the template parameter of the filter
32 // (defaults to float). For a 3D image, it will be organized as
33 // Datatype[z][y][x][c])
34 //
35 // The example also optionally takes in two more arguments, as a convenience to
36 // the user. These arguements extract the specified component 'c' from the
37 // membership image and rescale, so the user can fire up a typical image
38 // viewer and see the relative pixel memberships to class 'c'.
39 //
40 // Example args:
41 //   BrainProtonDensitySlice.png Memberships.mhd 4  2 Class2.png
42 //
43 // Here Memberships.mhd will be a 2x2x4 image containing pixel memberships
44 // Class2.png shows pixel memberships to the third class, (rescaled for display)
45 //
46 // Notes:
47 //   The default behaviour of the filter is to generate memberships by centering
48 // gaussian density functions around K-means of the pixel intensities in the
49 // image. The filter allows you to specify your own membership functions as well.
50 //
51 
52 #include "itkImage.h"
53 #include "itkBayesianClassifierInitializationImageFilter.h"
54 #include "itkImageFileReader.h"
55 #include "itkImageFileWriter.h"
56 #include "itkRescaleIntensityImageFilter.h"
57 #include "itkImageRegionConstIterator.h"
58 
main(int argc,char * argv[])59 int main(int argc, char *argv[])
60 {
61 
62   constexpr unsigned int Dimension = 2;
63   if( argc < 4 )
64     {
65     std::cerr << "Usage arguments: InputImage MembershipImage numberOfClasses [componentToExtract ExtractedImage]" << std::endl;
66     std::cerr << "  The MembershipImage image written is a VectorImage, ( an image with multiple components ) ";
67     std::cerr << "Given that most viewers can't see vector images, we will optionally extract a component and ";
68     std::cerr << "write it out as a scalar image as well." << std::endl;
69     return EXIT_FAILURE;
70     }
71 
72   using ImageType = itk::Image< unsigned char, Dimension >;
73   using BayesianInitializerType = itk::BayesianClassifierInitializationImageFilter<ImageType>;
74   BayesianInitializerType::Pointer bayesianInitializer
75                                           = BayesianInitializerType::New();
76 
77   using ReaderType = itk::ImageFileReader< ImageType >;
78   ReaderType::Pointer reader = ReaderType::New();
79   reader->SetFileName( argv[1] );
80 
81   try
82     {
83     reader->Update();
84     }
85   catch( itk::ExceptionObject & excp )
86     {
87     std::cerr << "Exception thrown " << std::endl;
88     std::cerr << excp << std::endl;
89     return EXIT_FAILURE;
90     }
91 
92   bayesianInitializer->SetInput( reader->GetOutput() );
93   bayesianInitializer->SetNumberOfClasses( std::stoi( argv[3] ) );
94 
95   // TODO add test where we specify membership functions
96 
97   using WriterType = itk::ImageFileWriter<BayesianInitializerType::OutputImageType>;
98   WriterType::Pointer writer = WriterType::New();
99   writer->SetInput( bayesianInitializer->GetOutput() );
100   writer->SetFileName( argv[2] );
101 
102   try
103     {
104     bayesianInitializer->Update();
105     }
106   catch( itk::ExceptionObject & excp )
107     {
108     std::cerr << "Exception thrown " << std::endl;
109     std::cerr << excp << std::endl;
110     return EXIT_FAILURE;
111     }
112 
113   try
114     {
115     writer->Update();
116     }
117   catch( itk::ExceptionObject & excp )
118     {
119     std::cerr << "Exception thrown " << std::endl;
120     std::cerr << excp << std::endl;
121     return EXIT_FAILURE;
122     }
123 
124   if( argv[4] && argv[5] )
125     {
126     using MembershipImageType = BayesianInitializerType::OutputImageType;
127     using ExtractedComponentImageType = itk::Image< MembershipImageType::InternalPixelType,
128                         Dimension >;
129     ExtractedComponentImageType::Pointer extractedComponentImage =
130                                     ExtractedComponentImageType::New();
131     extractedComponentImage->CopyInformation(
132                           bayesianInitializer->GetOutput() );
133     extractedComponentImage->SetBufferedRegion( bayesianInitializer->GetOutput()->GetBufferedRegion() );
134     extractedComponentImage->SetRequestedRegion( bayesianInitializer->GetOutput()->GetRequestedRegion() );
135     extractedComponentImage->Allocate();
136     using ConstIteratorType = itk::ImageRegionConstIterator< MembershipImageType >;
137     using IteratorType = itk::ImageRegionIterator< ExtractedComponentImageType >;
138     ConstIteratorType cit( bayesianInitializer->GetOutput(),
139                      bayesianInitializer->GetOutput()->GetBufferedRegion() );
140     IteratorType it( extractedComponentImage,
141                      extractedComponentImage->GetLargestPossibleRegion() );
142 
143     const unsigned int componentToExtract = std::stoi( argv[4] );
144     cit.GoToBegin();
145     it.GoToBegin();
146     while( !cit.IsAtEnd() )
147       {
148       it.Set(cit.Get()[componentToExtract]);
149       ++it;
150       ++cit;
151       }
152 
153     // Write out the rescaled extracted component
154     using OutputImageType = itk::Image< unsigned char, Dimension >;
155     using RescalerType = itk::RescaleIntensityImageFilter<
156       ExtractedComponentImageType, OutputImageType >;
157     RescalerType::Pointer rescaler = RescalerType::New();
158     rescaler->SetInput( extractedComponentImage );
159     rescaler->SetOutputMinimum( 0 );
160     rescaler->SetOutputMaximum( 255 );
161     using ExtractedComponentWriterType = itk::ImageFileWriter<OutputImageType>;
162     ExtractedComponentWriterType::Pointer
163                rescaledImageWriter = ExtractedComponentWriterType::New();
164     rescaledImageWriter->SetInput( rescaler->GetOutput() );
165     rescaledImageWriter->SetFileName( argv[5] );
166     rescaledImageWriter->Update();
167     }
168 
169   return EXIT_SUCCESS;
170 }
171