RainMaker.cpp

   1 /***********************************************************************
   2 RainMaker - Class to detect objects moving through a given range of
   3 depths in a depth image sequence to trigger rainfall on virtual terrain.
   4 Copyright (c) 2012-2015 Oliver Kreylos
   5
   6 This file is part of the Augmented Reality Sandbox (SARndbox).
   7
   8 The Augmented Reality Sandbox is free software; you can redistribute it
   9 and/or modify it under the terms of the GNU General Public License as
  10 published by the Free Software Foundation; either version 2 of the
  11 License, or (at your option) any later version.
  12
  13 The Augmented Reality Sandbox is distributed in the hope that it will be
  14 useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16 General Public License for more details.
  17
  18 You should have received a copy of the GNU General Public License along
  19 with the Augmented Reality Sandbox; if not, write to the Free Software
  20 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
  21 ***********************************************************************/
  22
  23 #include "RainMaker.h"
  24
  25 #include <Misc/FunctionCalls.h>
  26 #include <Geometry/HVector.h>
  27 #include <Geometry/Plane.h>
  28
  29 #include "FindBlobs.h"
  30
  31 template <>
  32 class BlobProperty<unsigned short> { // Class to calculate the 3D centroid of a blob in depth image space
  33     /* Elements: */
  34   private:
  35     double pxs, pys, pzs; // Accumulated components of centroid
  36     size_t numPixels; // Number of accumulated pixels
  37
  38     /* Constructors and destructors: */
  39   public:
  40     BlobProperty(void)
  41         : pxs(0.0), pys(0.0), pzs(0.0),
  42           numPixels(0) {
  43     }
  44
  45     /* Methods: */
  46     void addPixel(unsigned int x, unsigned int y, const unsigned short& pixelValue) {
  47         pxs += double(x);
  48         pys += double(y);
  49         pzs += double(pixelValue);
  50         ++numPixels;
  51     }
  52     void merge(const BlobProperty& other) {
  53         pxs += other.pxs;
  54         pys += other.pys;
  55         pzs += other.pzs;
  56         numPixels += other.numPixels;
  57     }
  58     size_t getNumPixels(void) const {
  59         return numPixels;
  60     }
  61     Geometry::Point<double, 3> calcCentroid(void)
  62     const { // Returns the centroid of the blob in depth image space
  63         return Geometry::Point<double, 3>(pxs / double(numPixels), pys / double(numPixels),
  64                                           pzs / double(numPixels));
  65     }
  66 };
  67
  68 template <>
  69 class BlobProperty<float> { // Class to calculate the 3D centroid of a blob in depth image space
  70     /* Elements: */
  71   private:
  72     double pxs, pys, pzs; // Accumulated components of centroid
  73     size_t numPixels; // Number of accumulated pixels
  74
  75     /* Constructors and destructors: */
  76   public:
  77     BlobProperty(void)
  78         : pxs(0.0), pys(0.0), pzs(0.0),
  79           numPixels(0) {
  80     }
  81
  82     /* Methods: */
  83     void addPixel(unsigned int x, unsigned int y, const float& pixelValue) {
  84         pxs += double(x);
  85         pys += double(y);
  86         pzs += double(pixelValue);
  87         ++numPixels;
  88     }
  89     void merge(const BlobProperty& other) {
  90         pxs += other.pxs;
  91         pys += other.pys;
  92         pzs += other.pzs;
  93         numPixels += other.numPixels;
  94     }
  95     size_t getNumPixels(void) const {
  96         return numPixels;
  97     }
  98     Geometry::Point<double, 3> calcCentroid(void)
  99     const { // Returns the centroid of the blob in depth image space
 100         return Geometry::Point<double, 3>(pxs / double(numPixels), pys / double(numPixels),
 101                                           pzs / double(numPixels));
 102     }
 103 };
 104
 105 class ValidPixelProperty { // Functor class to identify valid pixels in raw depth frames
 106     /* Elements: */
 107   private:
 108     float minPlane[4]; // Plane equation of the lower bound of valid depth values in depth image space
 109     float maxPlane[4]; // Plane equation of the upper bound of valid depth values in depth image space
 110     Geometry::Matrix<float, 3, 4>
 111     colorDepthHomography; // Homography from 3D depth image space into 2D color image space
 112     unsigned int colorSize[2]; // Width and height of color frames
 113     const unsigned char* colorFrame; // The current color frame
 114
 115     /* Constructors and destructors: */
 116   public:
 117     ValidPixelProperty(const float sMinPlane[4], const float sMaxPlane[4],
 118                        const Geometry::Matrix<float, 3, 4>& sColorDepthHomography, const unsigned int sColorSize[2])
 119         : colorDepthHomography(sColorDepthHomography),
 120           colorFrame(0) {
 121         /* Copy the min and max plane equations: */
 122         for(int i = 0; i < 4; ++i)
 123             minPlane[i] = sMinPlane[i];
 124         for(int i = 0; i < 4; ++i)
 125             maxPlane[i] = sMaxPlane[i];
 126
 127         /* Copy the color image size: */
 128         for(int i = 0; i < 2; ++i)
 129             colorSize[i] = sColorSize[i];
 130     }
 131
 132     /* Methods: */
 133   public:
 134     void setColorFrame(const unsigned char*
 135                        newColorFrame) { // Sets the color frame for the next blob extraction
 136         colorFrame = newColorFrame;
 137     }
 138     bool operator()(unsigned int x, unsigned int y, const unsigned short& pixel) const {
 139         return operator()(x, y, float(pixel));
 140     }
 141     bool operator()(unsigned int x, unsigned int y, const float& pixel) const {
 142         /* Plug the pixel into the plane equations to determine its validity: */
 143         float px = float(x) + 0.5f;
 144         float py = float(y) + 0.5f;
 145         float pz = pixel;
 146         float minD = minPlane[0] * px + minPlane[1] * py + minPlane[2] * pz + minPlane[3];
 147         float maxD = maxPlane[0] * px + maxPlane[1] * py + maxPlane[2] * pz + maxPlane[3];
 148         if(minD < 0.0f || maxD > 0.0f)
 149             return false;
 150
 151 #if 0
 152
 153         /* Project the pixel into the color frame: */
 154         Geometry::ComponentArray<float, 3> colorPos = colorDepthHomography *
 155                 Geometry::ComponentArray<float, 4>(px, py, pz, 1.0f);
 156         int cx = int(Math::floor(colorPos[0] / colorPos[2]));
 157         int cy = int(Math::floor(colorPos[1] / colorPos[2]));
 158         if(cx < 0 || cx >= colorSize[0] || cy < 0 || cy >= colorSize[1])
 159             return false;
 160
 161 #if 0
 162
 163         /* Check if the pixel is mostly black-ish: */
 164         const unsigned char* rgb = colorFrame + ((cy * colorSize[0] + cx) * 3);
 165         return rgb[0] < 64U && rgb[1] < 64U && rgb[2] < 64U;
 166
 167 #else
 168
 169         /* Normalize the pixel's color: */
 170         const unsigned char* rgb = colorFrame + ((cy * colorSize[0] + cx) * 3);
 171         unsigned char max = rgb[0];
 172         for(int i = 1; i < 3; ++i)
 173             if(max < rgb[i])
 174                 max = rgb[i];
 175         float rgb0[3];
 176         for(int i = 0; i < 3; ++i)
 177             rgb0[i] = float(rgb[i]) / float(max);
 178
 179         /* Check if the color is red-ish: */
 180         return rgb0[0] >= 0.8f && rgb0[1] < 0.25f && rgb0[2] < 0.25f;
 181
 182 #endif
 183
 184 #else
 185
 186         return true;
 187
 188 #endif
 189     }
 190 };
 191
 192 /**************************
 193 Methods of class RainMaker:
 194 **************************/
 195
 196 template <class DepthPixelParam>
 197 inline
 198 void RainMaker::extractBlobs(const Kinect::FrameBuffer& depthFrame, const ValidPixelProperty& vpp,
 199                              RainMaker::BlobList& blobsCc) {
 200     /* Extract raw blobs from the depth frame: */
 201     std::vector< ::Blob<DepthPixelParam> > blobsDic = findBlobs(depthSize,
 202             depthFrame.getData<DepthPixelParam>(), vpp);
 203
 204     /* Transform all blobs larger than the threshold to camera space: */
 205     blobsCc.reserve(blobsDic.size());
 206     for(typename std::vector< ::Blob<DepthPixelParam> >::const_iterator bIt = blobsDic.begin();
 207             bIt != blobsDic.end(); ++bIt)
 208         if(bIt->max[0] - bIt->min[0] >= minBlobSize && bIt->max[1] - bIt->min[1] >= minBlobSize) {
 209             Blob blobCc;
 210             Point centroidDic = bIt->blobProperty.calcCentroid();
 211             blobCc.centroid = depthProjection.transform(centroidDic);
 212
 213             /* Estimate the radius of the blob in camera space (this is admittedly ad-hoc): */
 214             double radiusDic = double(bIt->max[0] - bIt->min[0]) * 0.5;
 215             if(radiusDic > (bIt->max[1] - bIt->min[1]) * 0.5) {
 216                 radiusDic = (bIt->max[1] - bIt->min[1]) * 0.5;
 217                 blobCc.radius = Geometry::dist(depthProjection.transform(Point(centroidDic[0],
 218                                                centroidDic[1] + radiusDic, centroidDic[2])), blobCc.centroid);
 219             } else
 220                 blobCc.radius = Geometry::dist(depthProjection.transform(Point(centroidDic[0] + radiusDic,
 221                                                centroidDic[1], centroidDic[2])), blobCc.centroid);
 222
 223             /* Store the blob: */
 224             blobsCc.push_back(blobCc);
 225         }
 226 }
 227
 228 void* RainMaker::detectionThreadMethod(void) {
 229     unsigned int lastInputDepthFrameVersion = 0;
 230     unsigned int lastInputColorFrameVersion = 0;
 231
 232     /* Create a pixel validity decider: */
 233     ValidPixelProperty vpp(minPlane, maxPlane, colorDepthHomography, colorSize);
 234
 235     while(true) {
 236         Kinect::FrameBuffer depthFrame, colorFrame;
 237         {
 238             Threads::MutexCond::Lock inputLock(inputCond);
 239
 240             /* Wait until a new depth and color frame arrive, or the program shuts down: */
 241             while(runDetectionThread && (lastInputDepthFrameVersion == inputDepthFrameVersion
 242                                          || lastInputColorFrameVersion == inputColorFrameVersion))
 243                 inputCond.wait(inputLock);
 244
 245             /* Bail out if the program is shutting down: */
 246             if(!runDetectionThread)
 247                 break;
 248
 249             /* Work on the new frames: */
 250             depthFrame = inputDepthFrame;
 251             colorFrame = inputColorFrame;
 252             lastInputDepthFrameVersion = inputDepthFrameVersion;
 253             lastInputColorFrameVersion = inputColorFrameVersion;
 254         }
 255
 256         if(outputBlobsFunction != 0) {
 257             /* Set the most recent color frame in the pixel validator: */
 258             vpp.setColorFrame(colorFrame.getData<unsigned char>());
 259
 260             /* Detect all objects in the depth frame between the min and max planes: */
 261             BlobList blobsCc;
 262             if(depthIsFloat)
 263                 extractBlobs<float>(depthFrame, vpp, blobsCc);
 264             else
 265                 extractBlobs<unsigned short>(depthFrame, vpp, blobsCc);
 266
 267             /* Call the callback function: */
 268             (*outputBlobsFunction)(blobsCc);
 269         }
 270     }
 271
 272     return 0;
 273 }
 274
 275 RainMaker::RainMaker(const unsigned int sDepthSize[2], const unsigned int sColorSize[2],
 276                      const RainMaker::PTransform& sDepthProjection, const RainMaker::PTransform& sColorProjection,
 277                      const RainMaker::Plane& basePlane, double minElevation, double maxElevation, int sMinBlobSize)
 278     : depthIsFloat(false),
 279       outputBlobsFunction(0) {
 280     /* Remember the frame sizes: */
 281     for(int i = 0; i < 2; ++i)
 282         depthSize[i] = sDepthSize[i];
 283     for(int i = 0; i < 2; ++i)
 284         colorSize[i] = sColorSize[i];
 285
 286     /* Remember the depth and color projections: */
 287     depthProjection = sDepthProjection;
 288     colorProjection = sColorProjection;
 289
 290     /* Calculate the direct homography from depth image space to color image space: */
 291     PTransform hom = PTransform::scale(PTransform::Scale(double(colorSize[0]), double(colorSize[1]),
 292                                        1.0)); // Go to color image space
 293     hom *= colorProjection; // Go to color texture space
 294
 295     /* Remove the superfluous z component row: */
 296     for(int i = 0; i < 2; ++i)
 297         for(int j = 0; j < 4; ++j)
 298             colorDepthHomography(i, j) = float(hom.getMatrix()(i, j));
 299     for(int j = 0; j < 4; ++j)
 300         colorDepthHomography(2, j) = float(hom.getMatrix()(3, j));
 301
 302     /* Initialize the input frame slot: */
 303     inputDepthFrameVersion = 0;
 304     inputColorFrameVersion = 0;
 305
 306     /* Calculate the equations of the minimum and maximum elevation planes in camera space: */
 307     PTransform::HVector minPlaneCc(basePlane.getNormal());
 308     minPlaneCc[3] = -(basePlane.getOffset() + minElevation * basePlane.getNormal().mag());
 309     PTransform::HVector maxPlaneCc(basePlane.getNormal());
 310     maxPlaneCc[3] = -(basePlane.getOffset() + maxElevation * basePlane.getNormal().mag());
 311
 312     /* Transform the plane equations to depth image space and flip and swap the min and max planes because elevation increases opposite to raw depth: */
 313     PTransform::HVector minPlaneDic(depthProjection.getMatrix().transposeMultiply(minPlaneCc));
 314     double minPlaneScale = -1.0 / Geometry::mag(minPlaneDic.toVector());
 315     for(int i = 0; i < 4; ++i)
 316         maxPlane[i] = float(minPlaneDic[i] * minPlaneScale);
 317     PTransform::HVector maxPlaneDic(depthProjection.getMatrix().transposeMultiply(maxPlaneCc));
 318     double maxPlaneScale = -1.0 / Geometry::mag(maxPlaneDic.toVector());
 319     for(int i = 0; i < 4; ++i)
 320         minPlane[i] = float(maxPlaneDic[i] * maxPlaneScale);
 321
 322     /* Initialize the blob detector: */
 323     minBlobSize = sMinBlobSize;
 324
 325     /* Start the object detection thread: */
 326     runDetectionThread = true;
 327     detectionThread.start(this, &RainMaker::detectionThreadMethod);
 328 }
 329
 330 RainMaker::~RainMaker(void) {
 331     /* Shut down the object detection thread: */
 332     {
 333         Threads::MutexCond::Lock inputLock(inputCond);
 334         runDetectionThread = false;
 335         inputCond.signal();
 336     }
 337     detectionThread.join();
 338
 339     /* Release all allocated resources: */
 340     delete outputBlobsFunction;
 341 }
 342
 343 void RainMaker::setDepthIsFloat(bool newDepthIsFloat) {
 344     depthIsFloat = newDepthIsFloat;
 345 }
 346
 347 void RainMaker::setOutputBlobsFunction(RainMaker::OutputBlobsFunction* newOutputBlobsFunction) {
 348     delete outputBlobsFunction;
 349     outputBlobsFunction = newOutputBlobsFunction;
 350 }
 351
 352 void RainMaker::receiveRawDepthFrame(const Kinect::FrameBuffer& newDepthFrame) {
 353     Threads::MutexCond::Lock inputLock(inputCond);
 354
 355     /* Store the new buffer in the input buffer: */
 356     inputDepthFrame = newDepthFrame;
 357     ++inputDepthFrameVersion;
 358
 359     /* Signal the background thread: */
 360     inputCond.signal();
 361 }
 362
 363 void RainMaker::receiveRawColorFrame(const Kinect::FrameBuffer& newColorFrame) {
 364     Threads::MutexCond::Lock inputLock(inputCond);
 365
 366     /* Store the new buffer in the input buffer: */
 367     inputColorFrame = newColorFrame;
 368     ++inputColorFrameVersion;
 369
 370     /* Signal the background thread: */
 371     inputCond.signal();
 372 }