/** 
* @file SemSegTools.cpp
* @brief tools for semantic segmentation
* @author Erik Rodner
* @date 03/19/2009

*/
#include <iostream>

#include "SemSegTools.h"

using namespace OBJREC;

using namespace std;
using namespace NICE;

#undef DEBUG_LOCALIZATION

void SemSegTools::collectTrainingExamples ( 
			  const Config * conf,
			  const std::string & section,
			  const LabeledSet & train,
			  const ClassNames & cn,
			  Examples & examples,
			  vector<CachedExample *> & imgexamples )
{
    assert ( train.count() > 0 );
    examples.clear();
    imgexamples.clear();

    int grid_size_x = conf->gI(section, "grid_size_x", 5 );
    int grid_size_y = conf->gI(section, "grid_size_y", 5 );
    int grid_border_x = conf->gI(section, "grid_border_x", 20 );
    int grid_border_y = conf->gI(section, "grid_border_y", 20 );

    std::string selection = conf->gS(section, "train_selection" );

    set<int> classnoSelection;
    cn.getSelection ( selection, classnoSelection );
    
    bool useExcludedAsBG = conf->gB(section, "use_excluded_as_background", false );

    int backgroundClassNo = 0;
    
    if ( useExcludedAsBG ) 
    {
	backgroundClassNo = cn.classno("various");
	assert ( backgroundClassNo >= 0 );
    }

    LOOP_ALL_S (train)
    {
	EACH_INFO(image_classno,imgInfo);
	std::string imgfn = imgInfo.img();

	if ( ! imgInfo.hasLocalizationInfo() ) {
	    fprintf (stderr, "WARNING: NO localization info found for %s !\n",
		imgfn.c_str() );
	    continue;	
	}

	int xsize, ysize;
	CachedExample *ce = new CachedExample ( imgfn );
	ce->getImageSize ( xsize, ysize );
	imgexamples.push_back ( ce );

	const LocalizationResult *locResult = imgInfo.localization();
	if ( locResult->size() <= 0 ) {
	    fprintf (stderr, "WARNING: NO ground truth polygons found for %s !\n",
		imgfn.c_str());
	    continue;	
	}

	fprintf (stderr, "SemSegTools: Collecting pixel examples from localization info: %s\n", 
	    imgfn.c_str() );

	NICE::Image pixelLabels (xsize, ysize);
	pixelLabels.set(0);
	locResult->calcLabeledImage ( pixelLabels, cn.getBackgroundClass() );

#ifdef DEBUG_LOCALIZATION
	NICE::Image img (imgfn);
	showImage(img);
	showImage(pixelLabels);
#endif

	Example pce ( ce, 0, 0 );
	for ( int x = 0 ; x < xsize ; x += grid_size_x ) 
	    for ( int y = 0 ; y < ysize ; y += grid_size_y )  
	    {
		if ( (x >= grid_border_x) &&
		    ( y >= grid_border_y ) && ( x < xsize - grid_border_x ) &&
		    ( y < ysize - grid_border_x ) )
		{
		    pce.x = x; pce.y = y;
		    int classno = pixelLabels.getPixel(x,y);

		    if ( classnoSelection.find(classno) != classnoSelection.end() ) {
			examples.push_back ( pair<int, Example> (
			    classno, 
			    pce // FIXME: offset handling
			) );
		    } else if ( useExcludedAsBG ) {
			examples.push_back ( pair<int, Example> (
			    backgroundClassNo, 
			    pce // FIXME: offset handling
			) );
		    } 
		}
	    }
    }

    fprintf (stderr, "total number of examples: %d\n", (int)examples.size() );
}