std::optional Samples::FloorDetector::TryDetectFloorPlane()

in body-tracking-samples/floor_detector_sample/FloorDetector.cpp [128:215]
65 lines of code
13 McCabe index (conditional complexity)

std::optional<Samples::Plane> Samples::FloorDetector::TryDetectFloorPlane(
    const std::vector<k4a_float3_t>& cloudPoints,
    const k4a_imu_sample_t& imuSample,
    const k4a_calibration_t& sensorCalibration,
    size_t minimumFloorPointCount)
{
    auto gravity = TryEstimateGravityVectorForDepthCamera(imuSample, sensorCalibration);
    if (gravity.has_value() && !cloudPoints.empty())
    {
        // Up normal is opposite to gravity down vector.
        Samples::Vector up = (gravity.value() * -1).Normalized();

        // Compute elevation of cloud points (projections on the floor normal).
        std::vector<float> offsets(cloudPoints.size());
        for (size_t i = 0; i < cloudPoints.size(); ++i)
        {
            offsets[i] = up.Dot(Samples::Vector(cloudPoints[i]));
        }

        // There could be several horizontal planes in the scene (floor, tables, ceiling).
        // For the floor, look for lowest N points whose elevations are within a small range from each other.

        const float planeDisplacementRangeInMeters = 0.050f; // 5 cm in meters.
        const float planeMaxTiltInDeg = 5.0f;

        const int binAggregation = 6;
        assert(binAggregation >= 1);
        const float binSize = planeDisplacementRangeInMeters / binAggregation;
        const auto histBins = Histogram(offsets, binSize);
        
        // Cumulative histogram counts.
        std::vector<size_t> cumulativeHistCounts(histBins.size());
        cumulativeHistCounts[0] = histBins[0].count;
        for (size_t i = 1; i < histBins.size(); ++i)
        {
            cumulativeHistCounts[i] = cumulativeHistCounts[i - 1] + histBins[i].count;
        }

        assert(cumulativeHistCounts.back() == offsets.size());

        for (size_t i = 1; i + binAggregation < cumulativeHistCounts.size(); ++i)
        {
            size_t aggBinStart = i;                 // inclusive bin
            size_t aggBinEnd = i + binAggregation;  // exclusive bin
            size_t inlierCount = cumulativeHistCounts[aggBinEnd - 1] - cumulativeHistCounts[aggBinStart - 1];
            if (inlierCount > minimumFloorPointCount)
            {
                float offsetStart = histBins[aggBinStart].leftEdge; // inclusive
                float offsetEnd = histBins[aggBinEnd].leftEdge;     // exclusive

                // Inlier indices.
                std::vector<size_t> inlierIndices;
                for (size_t j = 0; j < offsets.size(); ++j)
                {
                    if (offsetStart <= offsets[j] && offsets[j] < offsetEnd)
                    {
                        inlierIndices.push_back(j);
                    }
                }

                // Fit plane to inlier points.
                auto refinedPlane = FitPlaneToInlierPoints(cloudPoints, inlierIndices);

                if (refinedPlane.has_value())
                {
                    // Ensure normal is upward.
                    if (refinedPlane->Normal.Dot(up) < 0)
                    {
                        refinedPlane->Normal = refinedPlane->Normal * -1;
                    }

                    // Ensure normal is mostly vertical.
                    auto floorTiltInDeg = acos(refinedPlane->Normal.Dot(up)) * 180.0f / 3.14159265f;
                    if (floorTiltInDeg < planeMaxTiltInDeg)
                    {
                        // For reduced jitter, use gravity for floor normal.
                        refinedPlane->Normal = up;
                        return refinedPlane;
                    }
                }

                return {};
            }
        }
    }

    return {};
}