NICE_SemSeg/progs/testSemanticSegmentation3D.cpp

// Beispielhafter Aufruf: BUILD_x86_64/progs/testSemanticSegmentation -config <CONFIGFILE>

/**
* @file testSemanticSegmentation.cpp
* @brief test semantic segmentation routines for 3d images and 2d images
* @author Erik Rodner, Björn Fröhlich, Sven Sickert
* @date 03/20/2008
*/

#ifdef NICE_USELIB_OPENMP
#include <omp.h>
#endif

#include "core/basics/Config.h"
#include "core/basics/StringTools.h"
#include "vislearning/baselib/ICETools.h"

#include "core/image/MultiChannelImage3DT.h"
#include "semseg/semseg/SemSegContextTree3D.h"
#include "semseg/semseg/SemSegTools.h"

#include "core/basics/ResourceStatistics.h"
#include "core/image/Morph.h"

#include <fstream>
#include <vector>

#undef DEBUG

using namespace OBJREC;

using namespace NICE;

using namespace std;

void startClassification (SemanticSegmentation *semseg,
                          std::vector< NICE::Matrix > & M_vec,
                          const Config & conf,
                          const LabeledSet* testFiles,
                          const ClassNames & classNames,
                          const set<int> & forbidden_classes,
                          map<int,int> & classMapping,
                          const string & resultdir,
                          const bool doCrossVal)
{
    bool show_results = conf.gB ( "debug", "show_results", false );
    bool write_results = conf.gB ( "debug", "write_results", false );
    bool writeProbMaps = conf.gB ( "debug", "write_prob_maps", false );
    if (doCrossVal)
        write_results = false;

    bool run_3Dseg = conf.gB( "SSContextTree", "run_3dseg", false);
    bool postProcessing = conf.gB( "main", "post_process", false);
    string output_type = conf.gS ( "debug", "output_type", "ppm" );
    string output_postfix = conf.gS ( "debug", "output_postfix", "" );

    vector< int > zsizeVec;
    semseg->getDepthVector ( testFiles, zsizeVec, run_3Dseg );
    int depthCount = 0, idx = 0;
    vector< string > filelist;
    NICE::MultiChannelImageT<int> segresult;
    NICE::MultiChannelImageT<int> gt;

    for (LabeledSet::const_iterator it = testFiles->begin(); it != testFiles->end(); it++)
    {
        for (std::vector<ImageInfo *>::const_iterator jt = it->second.begin();
             jt != it->second.end(); jt++)
        {
            ImageInfo & info = *(*jt);
            std::string file = info.img();
            filelist.push_back ( file );
            depthCount++;

            NICE::ImageT<int> lm;
            NICE::ImageT<int> lm_gt;
            if ( info.hasLocalizationInfo() )
            {
                const LocalizationResult *l_gt = info.localization();

                lm.resize ( l_gt->xsize, l_gt->ysize );
                lm.set ( 0 );

                lm_gt.resize ( l_gt->xsize, l_gt->ysize );
                lm_gt.set ( 0 );

                l_gt->calcLabeledImage ( lm, classNames.getBackgroundClass() );
#ifdef DEBUG
                cout << "testSemanticSegmentation3D: Generating Labeled NICE::Image (Ground-Truth)" << endl;
#endif
                l_gt->calcLabeledImage ( lm_gt, classNames.getBackgroundClass() );
            }
            segresult.addChannel ( lm );
            gt.addChannel ( lm_gt );

            int depthBoundary = 0;
            if ( run_3Dseg )
            {
                depthBoundary = zsizeVec[idx];
            }

            if ( depthCount < depthBoundary ) continue;

            NICE::MultiChannelImage3DT<double> probabilities;

            semseg->classify ( filelist, segresult, probabilities );

            // save to file
            for ( int z = 0; z < segresult.channels(); z++ )
            {
                std::string fname = StringTools::baseName ( filelist[z], false );

                if ( show_results || write_results )
                {
                    NICE::ColorImage orig ( filelist[z] );
                    NICE::ColorImage rgb;
                    NICE::ColorImage rgb_gt;
                    NICE::ColorImage ov_rgb;
                    NICE::ColorImage ov_rgb_gt;

                    for ( int y = 0 ; y < orig.height(); y++ )
                    {
                        for ( int x = 0 ; x < orig.width(); x++ )
                        {
                            lm.setPixel ( x, y, segresult.get ( x, y, ( uint ) z ) );
                            if ( run_3Dseg )
                                lm_gt.setPixel ( x, y, gt.get ( x, y, ( uint ) z ) );
                        }
                    }

                    // confusion matrix
                    NICE::Matrix M ( classMapping.size(), classMapping.size() );
                    M.set ( 0 );
                    SemSegTools::updateConfusionMatrix ( lm, lm_gt, M, forbidden_classes,
                                                         classMapping );
                    M_vec.push_back ( M );

                    classNames.labelToRGB ( lm, rgb );
                    classNames.labelToRGB ( lm_gt, rgb_gt );

                    if (postProcessing)
                    {
                        // median filter
                        for (int r = 0; r < 3; r++)
                        {
                            NICE::Image postIm(rgb.width(), rgb.height());
                            NICE::median(*(rgb.getChannel(r)), &postIm, 1);
                            for (int y = 0; y < rgb.height(); y++)
                                for (int x = 0; x < rgb.width(); x++)
                                    rgb.setPixel(x,y,r, postIm.getPixelQuick(x,y));
                        }
                    }

                    if ( write_results )
                    {
                        SemSegTools::segmentToOverlay ( orig.getChannel(1), rgb, ov_rgb );
                        SemSegTools::segmentToOverlay ( orig.getChannel(1), rgb_gt, ov_rgb_gt );

                        std::stringstream out;
                        if ( output_postfix.size() > 0 )
                            out << resultdir << "/" << fname << output_postfix;
                        else
                            out << resultdir << "/" << fname;
#ifdef DEBUG
                        cout << "Writing to file " << out.str() << "_*." << output_type << endl;
#endif
                        orig.write ( out.str() + "_orig." + output_type );
                        rgb.write ( out.str() + "_result." + output_type );
                        rgb_gt.write ( out.str() + "_groundtruth." + output_type );
                        ov_rgb.write ( out.str() + "_overlay_res." + output_type );
                        ov_rgb_gt.write ( out.str() + "_overlay_gt." + output_type );

                        // write Probability maps
                        if (writeProbMaps)
                        {
                            NICE::ColorImage prob_map( probabilities.width(), probabilities.height() );
                            prob_map.set(0,0,0);
                            int iNumChannels = probabilities.channels();
                            for ( int idxProbMap = 0; idxProbMap < iNumChannels; idxProbMap++)
                            {
                                for ( int y = 0 ; y < probabilities.height(); y++ )
                                {
                                    for ( int x = 0 ; x < probabilities.width(); x++ )
                                    {
                                        double probVal = probabilities.get( x, y, z, idxProbMap ) * 255.0;
                                        int tmp = round(probVal);
                                        for ( int c = 0 ; c < 3 ; c++ )
                                            prob_map.setPixel( x, y, c, tmp );
                                    }
                                }
                                std::stringstream ssFileProbMap;
                                //ssFileProbMap << out.str() << "_probs." << "c" << idxProbMap << "." << output_type;
                                ssFileProbMap << out.str() << "_probs." << "c-" << classNames.code( idxProbMap ) << "." << output_type;
                                //classNames
                                prob_map.write ( ssFileProbMap.str() );
                            }
                        }
                    }
                }
            }

            // prepare for new 3d image
            filelist.clear();
            segresult.reInit(0,0,0);
            gt.reInit(0,0,0);
            depthCount = 0;
            idx++;
        }
    }

    segresult.freeData();
}

/**
 test semantic segmentation routines
*/
int main ( int argc, char **argv )
{
    std::set_terminate ( __gnu_cxx::__verbose_terminate_handler );

    Config conf ( argc, argv );

    ResourceStatistics rs;

    /*---------------CONFIGURATION---------------*/
    bool doCrossVal = conf.gB ( "debug", "do_crossval", false );
    string resultdir = conf.gS ( "debug", "resultdir", "." );
    /*-------------------------------------------*/

#ifdef DEBUG
    cerr << "Writing Results to " << resultdir << endl;
#endif

    std::vector< NICE::Matrix > M_vec;

    MultiDataset md ( &conf );

    const ClassNames & classNames = md.getClassNames ( "train" );
    set<int> forbidden_classes;
    classNames.getSelection ( conf.gS ( "analysis", "forbidden_classes", "" ),
                              forbidden_classes );

    vector<bool> usedClasses ( classNames.numClasses(), true );
    for ( set<int>::const_iterator it = forbidden_classes.begin();
          it != forbidden_classes.end(); ++it)
    {
        usedClasses [ *it ] = false;
    }

    map<int,int> classMapping, classMappingInv;
    int j = 0;
    for ( int i = 0; i < usedClasses.size(); i++ )
        if (usedClasses[i])
        {
            classMapping[i] = j;
            classMappingInv[j] = i;
            j++;
        }


    // initialize semantic segmentation method
    SemanticSegmentation *semseg = NULL;

    // TRAINING AND TESTING
    if (!doCrossVal)
    {
        semseg = new SemSegContextTree3D ( &conf, &classNames );

        // STANDARD EVALUATION
        cout << "\nTRAINING" << endl;
        cout << "########\n" << endl;
        semseg->train( &md );

        cout << "\nCLASSIFICATION" << endl;
        cout << "##############\n" << endl;
        const LabeledSet *testFiles = md["test"];
        startClassification (semseg, M_vec, conf, testFiles, classNames,
                             forbidden_classes, classMapping, resultdir, doCrossVal );

        delete semseg;
    }
    else
    {
        // CROSS-VALIDATION
        for (int cval = 1; cval <= 10; cval++)
        {
            semseg = new SemSegContextTree3D ( &conf, &classNames );

            stringstream ss;
            ss << cval;
            string cvaltrain = "train_cv" + ss.str();
            string cvaltest = "test_cv" + ss.str();

            cout << "\nTRAINING " << cval << endl;
            cout << "###########\n" << endl;
            const LabeledSet *trainFiles = md[cvaltrain];
            semseg->train( trainFiles );

            cout << "\nCLASSIFICATION " << cval << endl;
            cout << "#################\n" << endl;
            const LabeledSet *testFiles = md[cvaltest];
            startClassification (semseg, M_vec, conf, testFiles, classNames,
                                 forbidden_classes, classMapping, resultdir, doCrossVal );

            delete semseg;
        }
    }

    cout << "\nSTATISTICS" << endl;
    cout << "##########\n" << endl;

    long maxMemory;
    double userCPUTime, sysCPUTime;
    rs.getStatistics ( maxMemory, userCPUTime, sysCPUTime );
    cout << "Memory (max):    " << maxMemory << " KB" << endl;
    cout << "CPU Time (user): " << userCPUTime << " seconds" << endl;
    cout << "CPU Time (sys):  " << sysCPUTime << " seconds" << endl;

    NICE::Matrix M ( classMapping.size(), classMapping.size() );
    M.set ( 0 );
    for ( int s = 0; s < ( int ) M_vec.size(); s++ )
    {
        NICE::Matrix M_tmp = M_vec[s];
        for ( int r = 0; r < ( int ) M_tmp.rows(); r++ )
            for ( int c = 0; c < ( int ) M_tmp.cols(); c++ )
                M ( r, c ) += M_tmp ( r, c );
    }
    // evaluation & analysis
    SemSegTools::computeClassificationStatistics(
                M, classNames, forbidden_classes, classMappingInv );

    return 0;
}