NICE_SemSeg/semseg/SemSegContextTree.cpp

#include "SemSegContextTree.h"
#include "vislearning/baselib/Globals.h"
#include "vislearning/baselib/ProgressBar.h"
#include "core/basics/StringTools.h"

#include "vislearning/cbaselib/CachedExample.h"
#include "vislearning/cbaselib/PascalResults.h"
#include "objrec/segmentation/RSMeanShift.h"
#include "objrec/segmentation/RSGraphBased.h"
#include "core/basics/numerictools.h"

#include <omp.h>

#include <iostream>

#define BOUND(x,min,max) (((x)<(min))?(min):((x)>(max)?(max):(x)))
#undef LOCALFEATS
//#define LOCALFEATS

using namespace OBJREC;
using namespace std;
using namespace NICE;

class Minus:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = feats.get(BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),channel1);
		double v2 = feats.get(BOUND(x+x2,0,xsize-1),BOUND(y+y2,0,ysize-1),channel2);
		return v1-v2;
	}
	
	virtual Operation* clone()
	{
		return new Minus();
	}
	
	virtual string writeInfos()
	{
		return "Minus";
	}
};

class MinusAbs:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = feats.get(BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),channel1);
		double v2 = feats.get(BOUND(x+x2,0,xsize-1),BOUND(y+y2,0,ysize-1),channel2);
		return abs(v1-v2);
	}
	
	virtual Operation* clone()
	{
		return new MinusAbs();
	};
	
	virtual string writeInfos()
	{
		return "MinusAbs";
	}
};

class Addition:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = feats.get(BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),channel1);
		double v2 = feats.get(BOUND(x+x2,0,xsize-1),BOUND(y+y2,0,ysize-1),channel2);
		return v1+v2;
	}
	
	virtual Operation* clone()
	{
		return new Addition();
	}
	
	virtual string writeInfos()
	{
		return "Addition";
	}
};

class Only1:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = feats.get(BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),channel1);
		return v1;
	}
	
	virtual Operation* clone()
	{
		return new Only1();
	}
	
	virtual string writeInfos()
	{
		return "Only1";
	}
};

class ContextMinus:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = tree[cfeats[BOUND(x+x1,0,xsize-1)][BOUND(y+y1,0,ysize-1)]].dist[channel1];
		double v2 = tree[cfeats[BOUND(x+x2,0,xsize-1)][BOUND(y+y2,0,ysize-1)]].dist[channel2];
		return v1-v2;
	}
	
	virtual Operation* clone()
	{
		return new ContextMinus();
	}
	
	virtual string writeInfos()
	{
		return "ContextMinus";
	}
};

class ContextMinusAbs:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = tree[cfeats[BOUND(x+x1,0,xsize-1)][BOUND(y+y1,0,ysize-1)]].dist[channel1];
		double v2 = tree[cfeats[BOUND(x+x2,0,xsize-1)][BOUND(y+y2,0,ysize-1)]].dist[channel2];
		return abs(v1-v2);
	}
	
	virtual Operation* clone()
	{
		return new ContextMinusAbs();
	}
	
	virtual string writeInfos()
	{
		return "ContextMinusAbs";
	}
};

class ContextAddition:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = tree[cfeats[BOUND(x+x1,0,xsize-1)][BOUND(y+y1,0,ysize-1)]].dist[channel1];
		double v2 = tree[cfeats[BOUND(x+x2,0,xsize-1)][BOUND(y+y2,0,ysize-1)]].dist[channel2];
		return v1+v2;
	}
	
	virtual Operation* clone()
	{
		return new ContextAddition();
	}
	
	virtual string writeInfos()
	{
		return "ContextAddition";
	}
};

class ContextOnly1:public Operation
{
public:
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		double v1 = tree[cfeats[BOUND(x+x1,0,xsize-1)][BOUND(y+y1,0,ysize-1)]].dist[channel1];
		return v1;
	}
	
	virtual Operation* clone()
	{
		return new ContextOnly1();
	}
	
	virtual string writeInfos()
	{
		return "ContextOnly1";
	}
};

// uses mean of classification in window given by (x1,y1) (x2,y2)
class IntegralOps:public Operation
{
public:
	virtual void set(int _x1, int _y1, int _x2, int _y2, int _channel1, int _channel2)
	{
		x1 = min(_x1,_x2);
		y1 = min(_y1,_y2);
		x2 = max(_x1,_x2);
		y2 = max(_y1,_y2);
		channel1 = _channel1;
		channel2 = _channel2;
	}
	
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		return computeMean(integralImg,BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),BOUND(x+x2,0,xsize-1),BOUND(y+y2,0,ysize-1),channel1);
	}
	
	inline double computeMean(const NICE::MultiChannelImageT<double> &intImg, const int &uLx, const int &uLy, const int &lRx, const int &lRy, const int &chan)
	{
		double val1 = intImg.get(uLx,uLy, chan);
		double val2 = intImg.get(lRx,uLy, chan);
		double val3 = intImg.get(uLx,lRy, chan);
		double val4 = intImg.get(lRx,lRy, chan);
		double area = (lRx-uLx)*(lRy-uLy);
		return (val1+val4-val2-val3)/area;
	}
	
	virtual Operation* clone()
	{
		return new IntegralOps();
	}
	
	virtual string writeInfos()
	{
		return "IntegralOps";
	}
};

//uses mean of Integral image given by x1, y1 with current pixel as center
class IntegralCenteredOps:public IntegralOps
{
public:
	virtual void set(int _x1, int _y1, int _x2, int _y2, int _channel1, int _channel2)
	{
		x1 = min(_x1,-_x1);
		y1 = min(_y1,-_y1);
		x2 = -x1;
		y2 = -y1;
		channel1 = _channel1;
		channel2 = _channel2;
	}
	
	virtual Operation* clone()
	{
		return new IntegralCenteredOps();
	}
	
	virtual string writeInfos()
	{
		return "IntegralCenteredOps";
	}
};

//uses different of mean of Integral image given by two windows, where (x1,y1) is the width and height of window1 and (x2,y2) of window 2
class BiIntegralCenteredOps:public IntegralOps
{
public:
	virtual void set(int _x1, int _y1, int _x2, int _y2, int _channel1, int _channel2)
	{
		x1 = min(abs(_x1),abs(_x2));
		y1 = min(abs(_y1),abs(_y2));
		x2 = max(abs(_x1),abs(_x2));
		y2 = max(abs(_y1),abs(_y2));
		channel1 = _channel1;
		channel2 = _channel2;
	}
	
	virtual double getVal(const NICE::MultiChannelImageT<double> &feats, const std::vector<std::vector<int> > &cfeats, const std::vector<TreeNode> &tree, MultiChannelImageT<double>  &integralImg, const int &x, const int &y)
	{
		int xsize = feats.width();
		int ysize = feats.height();
		return computeMean(integralImg,BOUND(x-x1,0,xsize-1),BOUND(y-y1,0,ysize-1),BOUND(x+x1,0,xsize-1),BOUND(y+y1,0,ysize-1),channel1) - computeMean(integralImg,BOUND(x-x2,0,xsize-1),BOUND(y-y2,0,ysize-1),BOUND(x+x2,0,xsize-1),BOUND(y+y2,0,ysize-1),channel1);
	}
	
	virtual Operation* clone()
	{
		return new BiIntegralCenteredOps();
	}
	
	virtual string writeInfos()
	{
		return "BiIntegralCenteredOps";
	}
};

SemSegContextTree::SemSegContextTree( const Config *conf, const MultiDataset *md )
    : SemanticSegmentation ( conf, &(md->getClassNames("train")) )
{
	this->conf = conf;
	string section = "SSContextTree";
	lfcw = new LFColorWeijer(conf);
	
	grid = conf->gI(section, "grid", 10 );
	
	maxSamples = conf->gI(section, "max_samples", 2000);
	
	minFeats = conf->gI(section, "min_feats", 50 );
	
	maxDepth = conf->gI(section, "max_depth", 10 );
	
	windowSize = conf->gI(section, "window_size", 16);
	
	featsPerSplit = conf->gI(section, "feats_per_split", 200);
	
	useShannonEntropy = conf->gB(section, "use_shannon_entropy", true);
	
	string segmentationtype = conf->gS(section, "segmentation_type", "meanshift");
	
	useGaussian = conf->gB(section, "use_gaussian", true);
	
	pixelWiseLabeling = false;
	
	if(segmentationtype == "meanshift")
		segmentation = new RSMeanShift(conf);
	else if (segmentationtype == "none")
	{
		segmentation = NULL;
		pixelWiseLabeling = true;
	}
	else if (segmentationtype == "felzenszwalb")
		segmentation = new RSGraphBased(conf);
	else
		throw("no valid segmenation_type\n please choose between none, meanshift and felzenszwalb\n");
	
	
	ftypes = conf->gI(section, "features", 2);;
	
	ops.push_back(new Minus());
	ops.push_back(new MinusAbs());
	ops.push_back(new Addition());
	ops.push_back(new Only1());
	
	cops.push_back(new ContextMinus());
	cops.push_back(new ContextMinusAbs());
	cops.push_back(new ContextAddition());
	cops.push_back(new ContextOnly1());
	cops.push_back(new BiIntegralCenteredOps());
	cops.push_back(new IntegralCenteredOps());
	cops.push_back(new IntegralOps());
	
	classnames = md->getClassNames ( "train" );
	
	///////////////////////////////////
	// Train Segmentation Context Trees
	///////////////////////////////////

	train ( md );
}

SemSegContextTree::~SemSegContextTree()
{
}

void SemSegContextTree::getBestSplit(const vector<MultiChannelImageT<double> > &feats, vector<vector<vector<int> > > &currentfeats, vector<MultiChannelImageT<double> > &integralImgs, const vector<vector<vector<int> > > &labels, int node, Operation *&splitop, double &splitval)
{

	int imgCount = 0, featdim = 0;
	try
	{
		imgCount = (int)feats.size();
		featdim = feats[0].channels();
	}
	catch(Exception)
	{
		cerr << "no features computed?" << endl;
	}

	double bestig = -numeric_limits< double >::max();
	splitop = NULL;
	splitval = -1.0;
	
	set<vector<int> >selFeats;
	map<int,int> e;
	int featcounter = 0;

	for(int iCounter = 0; iCounter < imgCount; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				if(currentfeats[iCounter][x][y] == node)
				{
					featcounter++;
				}
			}
		}
	}

	if(featcounter < minFeats)
	{
		cout << "only " << featcounter << " feats in current node -> it's a leaf" << endl;
		return;
	}
	
	vector<double> fraction(a.size(),0.0);
	for(uint i = 0; i < fraction.size(); i++)
	{
		if ( forbidden_classes.find ( labelmapback[i] ) != forbidden_classes.end() )
			fraction[i] = 0;
		else
			fraction[i] = ((double)maxSamples)/((double)featcounter*a[i]*a.size());
		//cout << "fraction["<<i<<"]: "<< fraction[i] << " a[" << i << "]: " << a[i] << endl;
	}
	//cout << "a.size(): " << a.size() << endl;
	//getchar();
	featcounter = 0;

	for(int iCounter = 0; iCounter < imgCount; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				if(currentfeats[iCounter][x][y] == node)
				{
					int cn = labels[iCounter][x][y];
					double randD = (double)rand()/(double)RAND_MAX;
					if(randD < fraction[labelmap[cn]])
					{
						vector<int> tmp(3,0);
						tmp[0] = iCounter;
						tmp[1] = x;
						tmp[2] = y;
						featcounter++;
						selFeats.insert(tmp);
						e[cn]++;
					}
				}
			}
		}
	}

	//cout << "size: " << selFeats.size() << endl;
	//getchar();
	
	map<int,int>::iterator mapit;
	double globent = 0.0;
	for ( mapit=e.begin() ; mapit != e.end(); mapit++ )
	{
	  //cout << "class: " << mapit->first << ": " << mapit->second << endl;
		double p = (double)(*mapit).second/(double)featcounter;
		globent += p*log2(p);
	}
	globent = -globent;
	
	if(globent < 0.5)
	{
		cout << "globent to small: " << globent << endl;
		return;
	}

	int classes = (int)labelmap.size();
	featsel.clear();
	for(int i = 0; i < featsPerSplit; i++)
	{
		int x1, x2, y1, y2;
		
		if(useGaussian)
		{
			double sigma = (double)windowSize/2.0;
			x1 = randGaussDouble(sigma)*(double)windowSize;
			x2 = randGaussDouble(sigma)*(double)windowSize;
			y1 = randGaussDouble(sigma)*(double)windowSize;
			y2 = randGaussDouble(sigma)*(double)windowSize;
		}
		else
		{
			x1 = (int)((double)rand()/(double)RAND_MAX*(double)windowSize)-windowSize/2;
			x2 = (int)((double)rand()/(double)RAND_MAX*(double)windowSize)-windowSize/2;
			y1 = (int)((double)rand()/(double)RAND_MAX*(double)windowSize)-windowSize/2;
			y2 = (int)((double)rand()/(double)RAND_MAX*(double)windowSize)-windowSize/2;
		}
		
		int ft = (int)((double)rand()/(double)RAND_MAX*(double)ftypes);

		if(integralImgs[0].width() == 0)
			ft = 0;
		
		if(ft == 0)
		{
			int f1 = (int)((double)rand()/(double)RAND_MAX*(double)featdim);
			int f2 = (int)((double)rand()/(double)RAND_MAX*(double)featdim);
			int o = (int)((double)rand()/(double)RAND_MAX*(double)ops.size());
			Operation *op = ops[o]->clone();
			op->set(x1,y1,x2,y2,f1,f2);
			featsel.push_back(op);
		}
		else if(ft == 1)
		{
			int f1 = (int)((double)rand()/(double)RAND_MAX*(double)classes);
			int f2 = (int)((double)rand()/(double)RAND_MAX*(double)classes);
			int o = (int)((double)rand()/(double)RAND_MAX*(double)cops.size());
			Operation *op = cops[o]->clone();
			op->set(x1,y1,x2,y2,f1,f2);
			featsel.push_back(op);
		}
	}
	
#pragma omp parallel for private(mapit)
	for(int f = 0; f < featsPerSplit; f++)
	{
		double l_bestig = -numeric_limits< double >::max();
		double l_splitval = -1.0;
		set<vector<int> >::iterator it;
		vector<double> vals;
		
		for ( it=selFeats.begin() ; it != selFeats.end(); it++ )
		{
			vals.push_back(featsel[f]->getVal(feats[(*it)[0]],currentfeats[(*it)[0]],tree, integralImgs[(*it)[0]], (*it)[1], (*it)[2]));
		}
		
		int counter = 0;
		for ( it=selFeats.begin() ; it != selFeats.end(); it++ , counter++)
		{
			set<vector<int> >::iterator it2;
			double val = vals[counter];
			
			map<int,int> eL, eR;
			int counterL = 0, counterR = 0;
			int counter2 = 0;
			for ( it2=selFeats.begin() ; it2 != selFeats.end(); it2++, counter2++ )
			{
				int cn = labels[(*it2)[0]][(*it2)[1]][(*it2)[2]];
				//cout << "vals[counter2] " << vals[counter2] << " val: " <<  val << endl;
				if(vals[counter2] < val)
				{
					//left entropie:
					eL[cn] = eL[cn]+1;
					counterL++;
				}
				else
				{
					//right entropie:
					eR[cn] = eR[cn]+1;
					counterR++;
				}
			}
			
			double leftent = 0.0;
			for ( mapit=eL.begin() ; mapit != eL.end(); mapit++ )
			{
				double p = (double)(*mapit).second/(double)counterL;
				leftent -= p*log2(p);
			}
			
			double rightent = 0.0;
			for ( mapit=eR.begin() ; mapit != eR.end(); mapit++ )
			{
				double p = (double)(*mapit).second/(double)counterR;
				rightent -= p*log2(p);
			}
			//cout << "rightent: " << rightent << " leftent: " << leftent << endl;
			
			double pl = (double)counterL/(double)(counterL+counterR);
			double ig = globent - (1.0-pl) * rightent - pl*leftent;
			//double ig = globent - rightent - leftent;
			
			if(useShannonEntropy)
			{
				double esplit = - ( pl*log(pl) + (1-pl)*log(1-pl) );
				ig = 2*ig / ( globent + esplit );
			}
			
			if(ig > l_bestig)
			{
				l_bestig = ig;
				l_splitval = val;
			}
		}
		
#pragma omp critical
{
		//cout << "globent: " << globent <<  " bestig " << bestig << " splitfeat: " << splitfeat << " splitval: " << splitval << endl;
		//cout << "globent: " << globent <<  " l_bestig " << l_bestig << " f: " << p << " l_splitval: " << l_splitval << endl;
		//cout << "p: " << featsubset[f] << endl;
		if(l_bestig > bestig)
		{
			bestig = l_bestig;
			splitop = featsel[f];
			splitval = l_splitval;
		}
}
	}
		//splitop->writeInfos();
		//cout<< "ig: " << bestig << endl;
	
	/*for(int i = 0; i < featsPerSplit; i++)
	{
		if(featsel[i] != splitop)
			delete featsel[i];
	}*/
#ifdef debug
	cout << "globent: " << globent <<  " bestig " << bestig << " splitval: " << splitval << endl;
#endif
}

void SemSegContextTree::computeIntegralImage(const vector<vector<int> > &currentfeats, MultiChannelImageT<double> &integralImage)
{
	int xsize = currentfeats.size();
	assert(xsize > 0);
	int ysize = currentfeats[0].size();
	
	int channels = (int)labelmap.size();
			
	for(int c = 0; c < channels; c++)
	{
		integralImage.set(0,0,tree[currentfeats[0][0]].dist[c], c);
		
		//first column
		for(int y = 1; y < ysize; y++)
		{
			integralImage.set(0,0,tree[currentfeats[0][y]].dist[c]+integralImage.get(0,y,c), c);
		}
		
		//first row
		for(int x = 1; x < xsize; x++)
		{
			integralImage.set(0,0,tree[currentfeats[x][0]].dist[c]+integralImage.get(x,0,c), c);
		}
		
		//rest
		for(int y = 1; y < ysize; y++)
		{
			for(int x = 1; x < xsize; x++)
			{
				double val = tree[currentfeats[x][y]].dist[c]+integralImage.get(x,y-1,c)+integralImage.get(x-1,y,c)-integralImage.get(x-1,y-1,c);
				integralImage.set(0, 0, val, c);
			}
		}
	}
}

void SemSegContextTree::train ( const MultiDataset *md )
{
	const LabeledSet train = * ( *md ) ["train"];
	const LabeledSet *trainp = &train;
	
	ProgressBar pb ( "compute feats" );
	pb.show();
	
	//TODO: Speichefresser!, lohnt sich sparse?
	vector<MultiChannelImageT<double> > allfeats;
	vector<vector<vector<int> > > currentfeats;
	vector<vector<vector<int> > > labels;

	std::string forbidden_classes_s = conf->gS ( "analysis", "donttrain", "" );
	if ( forbidden_classes_s == "" )
	{
		forbidden_classes_s = conf->gS ( "analysis", "forbidden_classes", "" );
	}
	
	classnames.getSelection ( forbidden_classes_s, forbidden_classes );
	
	int imgcounter = 0;
	
	LOOP_ALL_S ( *trainp )
	{
		EACH_INFO ( classno,info );

		NICE::ColorImage img;

		std::string currentFile = info.img();

		CachedExample *ce = new CachedExample ( currentFile );

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

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

		int xsize, ysize;
		ce->getImageSize ( xsize, ysize );

		vector<vector<int> > tmp = vector<vector<int> >(xsize, vector<int>(ysize,0));
		currentfeats.push_back(tmp);
		labels.push_back(tmp);

		try {
			img = ColorImage(currentFile);
		} catch (Exception) {
			cerr << "SemSeg: error opening image file <" << currentFile << ">" << endl;
			continue;
		}

		Globals::setCurrentImgFN ( currentFile );

		//TODO: resize image?!
		MultiChannelImageT<double> feats;
		allfeats.push_back(feats);
#ifdef LOCALFEATS
		lfcw->getFeats(img, allfeats[imgcounter]);
#else
		allfeats[imgcounter].reInit(xsize, ysize, 3, true);
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				for(int r = 0; r < 3; r++)
				{
					allfeats[imgcounter].set(x,y,img.getPixel(x,y,r),r);
				}
			}
		}
#endif		
		
		// getting groundtruth
		NICE::Image pixelLabels (xsize, ysize);
		pixelLabels.set(0);
		locResult->calcLabeledImage ( pixelLabels, ( *classNames ).getBackgroundClass() );
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				classno = pixelLabels.getPixel(x, y);
				labels[imgcounter][x][y] = classno;
				if ( forbidden_classes.find ( classno ) != forbidden_classes.end() )
					continue;
				labelcounter[classno]++;
				
			}
		}
		imgcounter++;
		pb.update ( trainp->count());
		delete ce;
	}
	pb.hide();
	
	map<int,int>::iterator mapit;
	int classes = 0;
	for(mapit = labelcounter.begin(); mapit != labelcounter.end(); mapit++)
	{
		labelmap[mapit->first] = classes;
		
		labelmapback[classes] = mapit->first;
		classes++;
	}
	
	//balancing
	int featcounter = 0;
	a = vector<double>(classes,0.0);
	for(int iCounter = 0; iCounter < imgcounter; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				featcounter++;
				int cn = labels[iCounter][x][y];
				a[labelmap[cn]] ++;
			}
		}
	}
	
	for(int i = 0; i < (int)a.size(); i++)
	{
		a[i] /= (double)featcounter;
	}
	
#ifdef DEBUG
	for(int i = 0; i < (int)a.size(); i++)
	{
		cout << "a["<<i<<"]: " << a[i] << endl;
	}
	cout << "a.size: " << a.size() << endl;
#endif
	
	tree.push_back(TreeNode());
	tree[0].dist = vector<double>(classes,0.0);
	int depth = 0;
	tree[0].depth = depth;
	int startnode = 0;
	bool allleaf = false;
	//int baseFeatSize = allfeats[0].size();
	
	vector<MultiChannelImageT<double> > integralImgs(imgcounter,MultiChannelImageT<double>());
	
	while(!allleaf && depth < maxDepth)
	{
		allleaf = true;
		//TODO vielleicht parallel wenn nächste schleife trotzdem noch parallelsiert würde, die hat mehr gewicht

		int t = (int) tree.size();
		int s = startnode;
		startnode = t;
		vector<vector<vector<int> > > lastfeats = currentfeats;
		
//#pragma omp parallel for
		for(int i = s; i < t; i++)
		{
			if(!tree[i].isleaf && tree[i].left < 0)
			{  
				Operation *splitfeat = NULL;
				double splitval;
				
				getBestSplit(allfeats, lastfeats, integralImgs, labels, i, splitfeat, splitval);
				tree[i].feat = splitfeat;
				
				tree[i].decision = splitval;
				
				if(splitfeat != NULL)
				{
					allleaf = false;
					int left = tree.size();
					tree.push_back(TreeNode());
					tree.push_back(TreeNode());
					int right = left+1;
					tree[i].left = left;
					tree[i].right = right;
					tree[left].dist = vector<double>(classes, 0.0);
					tree[right].dist = vector<double>(classes, 0.0);
					tree[left].depth = depth+1;
					tree[right].depth = depth+1;
#pragma omp parallel for
					for(int iCounter = 0; iCounter < imgcounter; iCounter++)
					{
						int xsize = currentfeats[iCounter].size();
						int ysize = currentfeats[iCounter][0].size();
						for(int x = 0; x < xsize; x++)
						{
							for(int y = 0; y < ysize; y++)
							{
								if(currentfeats[iCounter][x][y] == i)
								{
									double val = splitfeat->getVal(allfeats[iCounter],lastfeats[iCounter], tree, integralImgs[iCounter],x,y);
									if(val < splitval)
									{ 
										currentfeats[iCounter][x][y] = left;
										tree[left].dist[labelmap[labels[iCounter][x][y]]]++;
									}
									else
									{  
										currentfeats[iCounter][x][y] = right;
										tree[right].dist[labelmap[labels[iCounter][x][y]]]++;
									}
								}
							}
						}
					}
					
					double lcounter = 0.0, rcounter = 0.0;
					for(uint d = 0; d < tree[left].dist.size(); d++)
					{
						if ( forbidden_classes.find ( labelmapback[d] ) != forbidden_classes.end() )
						{
							tree[left].dist[d] = 0;
							tree[right].dist[d] = 0;
						}
						else
						{
							tree[left].dist[d]/=a[d];
							lcounter +=tree[left].dist[d];
							tree[right].dist[d]/=a[d];
							rcounter +=tree[right].dist[d];
						}
					}
					
					if(lcounter <= 0 || rcounter <= 0)
					{
						cout << "lcounter : " << lcounter << " rcounter: " << rcounter << endl;
						cout << "splitval: " << splitval << endl;
						assert(lcounter > 0 && rcounter > 0);
					}
					for(uint d = 0; d < tree[left].dist.size(); d++)
					{
						tree[left].dist[d]/=lcounter;
						tree[right].dist[d]/=rcounter;
					}
				}
				else
				{
					tree[i].isleaf = true;
				}
			}
		}
		//TODO: features neu berechnen!
		
		//compute integral image
		int channels = (int)labelmap.size();
		
		if(integralImgs[0].width() == 0)
		{
			for(int i = 0; i < imgcounter; i++)
			{
				int xsize = allfeats[i].width();
				int ysize = allfeats[i].height();
				
				integralImgs[i].reInit(xsize, ysize, channels);
			}
		}
		
		for(int i = 0; i < imgcounter; i++)
		{
			computeIntegralImage(currentfeats[i],integralImgs[i]);
		}
		
		if(depth == 4)
		{
			cout << "learn: ";
			for(int x = 20; x < 30; x++)
			{
				cout << currentfeats[0][x][20] << " ";
			}
			cout << endl;
		}
		
		depth++;
#ifdef DEBUG
		cout << "depth: " << depth << endl;
#endif
	}
	
	
#ifdef DEBUG
	int t = (int) tree.size();
	for(int i = 0; i < t; i++)
	{
		printf("tree[%i]: left: %i, right: %i", i, tree[i].left, tree[i].right);
		if(!tree[i].isleaf && tree[i].left != -1)
			cout <<  ", feat: " << tree[i].feat->writeInfos() << " ";
		for(int d = 0; d < (int)tree[i].dist.size(); d++)
		{
			cout << " " << tree[i].dist[d];
		}
		cout << endl;
	}
#endif
}

void SemSegContextTree::semanticseg ( CachedExample *ce, NICE::Image & segresult,NICE::MultiChannelImageT<double> & probabilities )
{
	int xsize;
	int ysize;
	ce->getImageSize ( xsize, ysize );
	
	int numClasses = classNames->numClasses();
	
	fprintf (stderr, "ContextTree classification !\n");

	probabilities.reInit ( xsize, ysize, numClasses, true );
	probabilities.setAll ( 0 );

	NICE::ColorImage img;

	std::string currentFile = Globals::getCurrentImgFN();
	
	try {
		img = ColorImage(currentFile);
	} catch (Exception) {
		cerr << "SemSeg: error opening image file <" << currentFile << ">" << endl;
		return;
	}
	
	//TODO: resize image?!
		
		MultiChannelImageT<double> feats;
#ifdef LOCALFEATS
		lfcw->getFeats(img, feats);
#else
		feats.reInit (xsize, ysize, 3, true);
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				for(int r = 0; r < 3; r++)
				{
					feats.set(x,y,img.getPixel(x,y,r),r);
				}
			}
		}
#endif
	
	bool allleaf = false;
	
	MultiChannelImageT<double> integralImg;
	
	vector<vector<int> > currentfeats = vector<vector<int> >(xsize, vector<int>(ysize,0));
	int depth = 0;
	while(!allleaf)
	{
		allleaf = true;
		//TODO vielleicht parallel wenn nächste schleife auch noch parallelsiert würde, die hat mehr gewicht
		//#pragma omp parallel for
		vector<vector<int> > lastfeats = currentfeats;
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				int t = currentfeats[x][y];
				if(tree[t].left > 0)
				{
					allleaf = false;
					double val = tree[t].feat->getVal(feats,lastfeats,tree,integralImg,x,y);
					
					if(val < tree[t].decision)
					{
						currentfeats[x][y] = tree[t].left;
					}
					else
					{
						currentfeats[x][y] = tree[t].right;
					}
				}
			}
		}
		
		//compute integral image
		int channels = (int)labelmap.size();
		
		if(integralImg.width() == 0)
		{
			int xsize = feats.width();
			int ysize = feats.height();
			
			integralImg.reInit(xsize, ysize, channels);
		}
		
		computeIntegralImage(currentfeats,integralImg);
		
		if(depth == 4)
		{
			cout << "learn: ";
			for(int x = 20; x < 30; x++)
			{
				cout << currentfeats[x][20] << " ";
			}
			cout << endl;
		}
		
		depth++;
	}
	
	if(pixelWiseLabeling)
	{
		//finales labeln:
		long int offset = 0;
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++,offset++)
			{
				int t = currentfeats[x][y];
				double maxvalue = - numeric_limits<double>::max(); //TODO: das muss nur pro knoten gemacht werden, nicht pro pixel
				int maxindex = 0;
				for(uint i = 0; i < tree[i].dist.size(); i++)
				{
					probabilities.data[labelmapback[i]][offset] = tree[t].dist[i];
					if(tree[t].dist[i] > maxvalue)
					{
						maxvalue = tree[t].dist[i];
						maxindex = labelmapback[i];
					}
					segresult.setPixel(x,y,maxindex);
				}
			}
		}
	}
	else
	{
		//final labeling using segmentation
		//TODO: segmentation
		Matrix regions;
		int regionNumber = segmentation->segRegions(img,regions);
		cout << "regions: " << regionNumber << endl;
		int dSize = (int)labelmap.size();
		vector<vector<double> > regionProbs(regionNumber, vector<double>(dSize,0.0));
		vector<int> bestlabels(regionNumber, 0);
		
		for(int y = 0; y < img.height(); y++)
		{
			for(int x = 0; x < img.width(); x++)
			{
				int cnode = currentfeats[x][y];
				int cregion = regions(x,y);
				for(int d = 0; d < dSize; d++)
				{
					regionProbs[cregion][d]+=tree[cnode].dist[d];
				}
			}
		}
		
		for(int r = 0; r < regionNumber; r++)
		{
			double maxval = regionProbs[r][0];
			for(int d = 1; d < dSize; d++)
			{
				if(maxval < regionProbs[r][d])
				{
					maxval = regionProbs[r][d];
					bestlabels[r] = d;
				}
			}
			bestlabels[r] = labelmapback[bestlabels[r]];
		}
		
		for(int y = 0; y < img.height(); y++)
		{
			for(int x = 0; x < img.width(); x++)
			{
				segresult.setPixel(x,y,bestlabels[regions(x,y)]);
			}
		}
	}
}