RainMaker.cpp

/***********************************************************************
RainMaker - Class to detect objects moving through a given range of
depths in a depth image sequence to trigger rainfall on virtual terrain.
Copyright (c) 2012-2013 Oliver Kreylos

This file is part of the Augmented Reality Sandbox (SARndbox).

The Augmented Reality Sandbox is free software; you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of the
License, or (at your option) any later version.

The Augmented Reality Sandbox is distributed in the hope that it will be
useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License along
with the Augmented Reality Sandbox; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
***********************************************************************/

#include "RainMaker.h"

#include <Misc/FunctionCalls.h>
#include <Geometry/HVector.h>
#include <Geometry/Plane.h>

#include "FindBlobs.h"

template <>
class BlobProperty<unsigned short> // Class to calculate the 3D centroid of a blob in depth image space
	{
	/* Elements: */
	private:
	double pxs,pys,pzs; // Accumulated components of centroid
	size_t numPixels; // Number of accumulated pixels
	
	/* Constructors and destructors: */
	public:
	BlobProperty(void)
		:pxs(0.0),pys(0.0),pzs(0.0),
		 numPixels(0)
		{
		}
	
	/* Methods: */
	void addPixel(unsigned int x,unsigned int y,const unsigned short& pixelValue)
		{
		pxs+=double(x);
		pys+=double(y);
		pzs+=double(pixelValue);
		++numPixels;
		}
	void merge(const BlobProperty& other)
		{
		pxs+=other.pxs;
		pys+=other.pys;
		pzs+=other.pzs;
		numPixels+=other.numPixels;
		}
	size_t getNumPixels(void) const
		{
		return numPixels;
		}
	Geometry::Point<double,3> calcCentroid(void) const // Returns the centroid of the blob in depth image space
		{
		return Geometry::Point<double,3>(pxs/double(numPixels),pys/double(numPixels),pzs/double(numPixels));
		}
	};

template <>
class BlobProperty<float> // Class to calculate the 3D centroid of a blob in depth image space
	{
	/* Elements: */
	private:
	double pxs,pys,pzs; // Accumulated components of centroid
	size_t numPixels; // Number of accumulated pixels
	
	/* Constructors and destructors: */
	public:
	BlobProperty(void)
		:pxs(0.0),pys(0.0),pzs(0.0),
		 numPixels(0)
		{
		}
	
	/* Methods: */
	void addPixel(unsigned int x,unsigned int y,const float& pixelValue)
		{
		pxs+=double(x);
		pys+=double(y);
		pzs+=double(pixelValue);
		++numPixels;
		}
	void merge(const BlobProperty& other)
		{
		pxs+=other.pxs;
		pys+=other.pys;
		pzs+=other.pzs;
		numPixels+=other.numPixels;
		}
	size_t getNumPixels(void) const
		{
		return numPixels;
		}
	Geometry::Point<double,3> calcCentroid(void) const // Returns the centroid of the blob in depth image space
		{
		return Geometry::Point<double,3>(pxs/double(numPixels),pys/double(numPixels),pzs/double(numPixels));
		}
	};

class ValidPixelProperty // Functor class to identify valid pixels in raw depth frames
	{
	/* Elements: */
	private:
	float minPlane[4]; // Plane equation of the lower bound of valid depth values in depth image space
	float maxPlane[4]; // Plane equation of the upper bound of valid depth values in depth image space
	Geometry::Matrix<float,3,4> colorDepthHomography; // Homography from 3D depth image space into 2D color image space
	unsigned int colorSize[2]; // Width and height of color frames
	const unsigned char* colorFrame; // The current color frame
	
	/* Constructors and destructors: */
	public:
	ValidPixelProperty(const float sMinPlane[4],const float sMaxPlane[4],const Geometry::Matrix<float,3,4>& sColorDepthHomography,const unsigned int sColorSize[2])
		:colorDepthHomography(sColorDepthHomography),
		 colorFrame(0)
		{
		/* Copy the min and max plane equations: */
		for(int i=0;i<4;++i)
			minPlane[i]=sMinPlane[i];
		for(int i=0;i<4;++i)
			maxPlane[i]=sMaxPlane[i];
		
		/* Copy the color image size: */
		for(int i=0;i<2;++i)
			colorSize[i]=sColorSize[i];
		}
	
	/* Methods: */
	public:
	void setColorFrame(const unsigned char* newColorFrame) // Sets the color frame for the next blob extraction
		{
		colorFrame=newColorFrame;
		}
	bool operator()(unsigned int x,unsigned int y,const unsigned short& pixel) const
		{
		return operator()(x,y,float(pixel));
		}
	bool operator()(unsigned int x,unsigned int y,const float& pixel) const
		{
		/* Plug the pixel into the plane equations to determine its validity: */
		float px=float(x)+0.5f;
		float py=float(y)+0.5f;
		float pz=pixel;
		float minD=minPlane[0]*px+minPlane[1]*py+minPlane[2]*pz+minPlane[3];
		float maxD=maxPlane[0]*px+maxPlane[1]*py+maxPlane[2]*pz+maxPlane[3];
		if(minD<0.0f||maxD>0.0f)
			return false;
		
		#if 0
		
		/* Project the pixel into the color frame: */
		Geometry::ComponentArray<float,3> colorPos=colorDepthHomography*Geometry::ComponentArray<float,4>(px,py,pz,1.0f);
		int cx=int(Math::floor(colorPos[0]/colorPos[2]));
		int cy=int(Math::floor(colorPos[1]/colorPos[2]));
		if(cx<0||cx>=colorSize[0]||cy<0||cy>=colorSize[1])
			return false;
		
		#if 0
		
		/* Check if the pixel is mostly black-ish: */
		const unsigned char* rgb=colorFrame+((cy*colorSize[0]+cx)*3);
		return rgb[0]<64U&&rgb[1]<64U&&rgb[2]<64U;
		
		#else
		
		/* Normalize the pixel's color: */
		const unsigned char* rgb=colorFrame+((cy*colorSize[0]+cx)*3);
		unsigned char max=rgb[0];
		for(int i=1;i<3;++i)
			if(max<rgb[i])
				max=rgb[i];
		float rgb0[3];
		for(int i=0;i<3;++i)
			rgb0[i]=float(rgb[i])/float(max);
		
		/* Check if the color is red-ish: */
		return rgb0[0]>=0.8f&&rgb0[1]<0.25f&&rgb0[2]<0.25f;
		
		#endif
		
		#else
		
		return true;
		
		#endif
		}
	};

/**************************
Methods of class RainMaker:
**************************/

template <class DepthPixelParam>
inline
void RainMaker::extractBlobs(const Kinect::FrameBuffer& depthFrame,const ValidPixelProperty& vpp,RainMaker::BlobList& blobsCc)
	{
	/* Extract raw blobs from the depth frame: */
	std::vector< ::Blob<DepthPixelParam> > blobsDic=findBlobs(depthSize,static_cast<const DepthPixelParam*>(depthFrame.getBuffer()),vpp);
	
	/* Transform all blobs larger than the threshold to camera space: */
	blobsCc.reserve(blobsDic.size());
	for(typename std::vector< ::Blob<DepthPixelParam> >::const_iterator bIt=blobsDic.begin();bIt!=blobsDic.end();++bIt)
		if(bIt->max[0]-bIt->min[0]>=minBlobSize&&bIt->max[1]-bIt->min[1]>=minBlobSize)
			{
			Blob blobCc;
			Point centroidDic=bIt->blobProperty.calcCentroid();
			blobCc.centroid=depthProjection.transform(centroidDic);
			
			/* Estimate the radius of the blob in camera space (this is admittedly ad-hoc): */
			double radiusDic=double(bIt->max[0]-bIt->min[0])*0.5;
			if(radiusDic>(bIt->max[1]-bIt->min[1])*0.5)
				{
				radiusDic=(bIt->max[1]-bIt->min[1])*0.5;
				blobCc.radius=Geometry::dist(depthProjection.transform(Point(centroidDic[0],centroidDic[1]+radiusDic,centroidDic[2])),blobCc.centroid);
				}
			else
				blobCc.radius=Geometry::dist(depthProjection.transform(Point(centroidDic[0]+radiusDic,centroidDic[1],centroidDic[2])),blobCc.centroid);
			
			/* Store the blob: */
			blobsCc.push_back(blobCc);
			}
	}

void* RainMaker::detectionThreadMethod(void)
	{
	unsigned int lastInputDepthFrameVersion=0;
	unsigned int lastInputColorFrameVersion=0;
	
	/* Create a pixel validity decider: */
	ValidPixelProperty vpp(minPlane,maxPlane,colorDepthHomography,colorSize);
	
	while(true)
		{
		Kinect::FrameBuffer depthFrame,colorFrame;
		{
		Threads::MutexCond::Lock inputLock(inputCond);
		
		/* Wait until a new depth and color frame arrive, or the program shuts down: */
		while(runDetectionThread&&(lastInputDepthFrameVersion==inputDepthFrameVersion||lastInputColorFrameVersion==inputColorFrameVersion))
			inputCond.wait(inputLock);
		
		/* Bail out if the program is shutting down: */
		if(!runDetectionThread)
			break;
		
		/* Work on the new frames: */
		depthFrame=inputDepthFrame;
		colorFrame=inputColorFrame;
		lastInputDepthFrameVersion=inputDepthFrameVersion;
		lastInputColorFrameVersion=inputColorFrameVersion;
		}
		
		if(outputBlobsFunction!=0)
			{
			/* Set the most recent color frame in the pixel validator: */
			vpp.setColorFrame(static_cast<const unsigned char*>(colorFrame.getBuffer()));
			
			/* Detect all objects in the depth frame between the min and max planes: */
			BlobList blobsCc;
			if(depthIsFloat)
				extractBlobs<float>(depthFrame,vpp,blobsCc);
			else
				extractBlobs<unsigned short>(depthFrame,vpp,blobsCc);
			
			/* Call the callback function: */
			(*outputBlobsFunction)(blobsCc);
			}
		}
	
	return 0;
	}

RainMaker::RainMaker(const unsigned int sDepthSize[2],const unsigned int sColorSize[2],const RainMaker::PTransform& sDepthProjection,const RainMaker::PTransform& sColorProjection,const RainMaker::Plane& basePlane,double minElevation,double maxElevation,int sMinBlobSize)
	:depthIsFloat(false),
	 outputBlobsFunction(0)
	{
	/* Remember the frame sizes: */
	for(int i=0;i<2;++i)
		depthSize[i]=sDepthSize[i];
	for(int i=0;i<2;++i)
		colorSize[i]=sColorSize[i];
	
	/* Remember the depth and color projections: */
	depthProjection=sDepthProjection;
	colorProjection=sColorProjection;
	
	/* Calculate the direct homography from depth image space to color image space: */
	PTransform hom=PTransform::scale(PTransform::Scale(double(colorSize[0]),double(colorSize[1]),1.0)); // Go to color image space
	hom*=colorProjection; // Go to color texture space
	
	/* Remove the superfluous z component row: */
	for(int i=0;i<2;++i)
		for(int j=0;j<4;++j)
			colorDepthHomography(i,j)=float(hom.getMatrix()(i,j));
	for(int j=0;j<4;++j)
		colorDepthHomography(2,j)=float(hom.getMatrix()(3,j));
	
	/* Initialize the input frame slot: */
	inputDepthFrameVersion=0;
	inputColorFrameVersion=0;
	
	/* Calculate the equations of the minimum and maximum elevation planes in camera space: */
	PTransform::HVector minPlaneCc(basePlane.getNormal());
	minPlaneCc[3]=-(basePlane.getOffset()+minElevation*basePlane.getNormal().mag());
	PTransform::HVector maxPlaneCc(basePlane.getNormal());
	maxPlaneCc[3]=-(basePlane.getOffset()+maxElevation*basePlane.getNormal().mag());
	
	/* Transform the plane equations to depth image space and flip and swap the min and max planes because elevation increases opposite to raw depth: */
	PTransform::HVector minPlaneDic(depthProjection.getMatrix().transposeMultiply(minPlaneCc));
	double minPlaneScale=-1.0/Geometry::mag(minPlaneDic.toVector());
	for(int i=0;i<4;++i)
		maxPlane[i]=float(minPlaneDic[i]*minPlaneScale);
	PTransform::HVector maxPlaneDic(depthProjection.getMatrix().transposeMultiply(maxPlaneCc));
	double maxPlaneScale=-1.0/Geometry::mag(maxPlaneDic.toVector());
	for(int i=0;i<4;++i)
		minPlane[i]=float(maxPlaneDic[i]*maxPlaneScale);
	
	/* Initialize the blob detector: */
	minBlobSize=sMinBlobSize;
	
	/* Start the object detection thread: */
	runDetectionThread=true;
	detectionThread.start(this,&RainMaker::detectionThreadMethod);
	}

RainMaker::~RainMaker(void)
	{
	/* Shut down the object detection thread: */
	{
	Threads::MutexCond::Lock inputLock(inputCond);
	runDetectionThread=false;
	inputCond.signal();
	}
	detectionThread.join();
	
	/* Release all allocated resources: */
	delete outputBlobsFunction;
	}

void RainMaker::setDepthIsFloat(bool newDepthIsFloat)
	{
	depthIsFloat=newDepthIsFloat;
	}

void RainMaker::setOutputBlobsFunction(RainMaker::OutputBlobsFunction* newOutputBlobsFunction)
	{
	delete outputBlobsFunction;
	outputBlobsFunction=newOutputBlobsFunction;
	}

void RainMaker::receiveRawDepthFrame(const Kinect::FrameBuffer& newDepthFrame)
	{
	Threads::MutexCond::Lock inputLock(inputCond);
	
	/* Store the new buffer in the input buffer: */
	inputDepthFrame=newDepthFrame;
	++inputDepthFrameVersion;
	
	/* Signal the background thread: */
	inputCond.signal();
	}

void RainMaker::receiveRawColorFrame(const Kinect::FrameBuffer& newColorFrame)
	{
	Threads::MutexCond::Lock inputLock(inputCond);
	
	/* Store the new buffer in the input buffer: */
	inputColorFrame=newColorFrame;
	++inputColorFrameVersion;
	
	/* Signal the background thread: */
	inputCond.signal();
	}