preprocessing_foreground_container_impl.hpp

// preprocessing_foreground_container_impl.hpp
// Created by greg on 1/13/25.
// Copyright (C) 2025 Presage Security, Inc.
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 3 of the License, or (at your option) any later version.
//
// This program 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
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program; if not, write to the Free Software Foundation,
// Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
// === standard library includes (if any) ===
// === third-party includes (if any) ===
#include <mediapipe/framework/port/status_macros.h>
#include <mediapipe/framework/port/opencv_core_inc.h>
#include <mediapipe/framework/port/opencv_imgproc_inc.h>
#include <physiology/graph/stream_and_packet_names.h>
#include <physiology/modules/messages/metrics.h>
// === local includes (if any) ===
#include "preprocessing_foreground_container.hpp"
#include "preprocessing_stream_and_packet_names.hpp"
#include <smartspectra/container/packet_helpers.hpp>
 
 
#pragma once
 
namespace presage::smartspectra::lab {
namespace c = container;
namespace poller = container::output_stream_poller_wrapper;
namespace pe = physiology::edge;
namespace pg = preprocessing_graph;
namespace ph = container::packet_helpers;
 
 
template<
    platform_independence::DeviceType TDeviceType,
    c::settings::OperationMode TOperationMode,
    c::settings::IntegrationMode TIntegrationMode
>
absl::Status PreprocessingForegroundContainer<TDeviceType,
    TOperationMode,
    TIntegrationMode>::InitializeOutputDataPollers() {
    MP_RETURN_IF_ERROR(this->preprocessed_data_poller.Initialize(this->graph, pg::output_streams::kPreprocessedData));
    // Initialize optional debug pollers (safe even if not used)
    // These streams exist in preprocessing graphs when tracking is enabled.
    // Ignore Initialize errors to avoid failing runs in graphs missing these outputs.
    (void)this->points_chest_poller.Initialize(this->graph, pg::output_streams::kPointsChest);
    (void)this->points_abdomen_poller.Initialize(this->graph, pg::output_streams::kPointsAbdomen);
    (void)this->points_glutes_poller.Initialize(this->graph, pg::output_streams::kPointsGlutes);
    (void)this->points_knees_poller.Initialize(this->graph, pg::output_streams::kPointsKnees);
    (void)this->transformed_size_poller.Initialize(this->graph, pg::output_streams::kTransformedSize);
    return absl::OkStatus();
}
 
template<platform_independence::DeviceType TDeviceType, container::settings::OperationMode TOperationMode, container::settings::IntegrationMode TIntegrationMode>
PreprocessingForegroundContainer<TDeviceType, TOperationMode, TIntegrationMode>::PreprocessingForegroundContainer(
    PreprocessingForegroundContainer::SettingsType settings
): Base(settings) {}
 
template<
    platform_independence::DeviceType TDeviceType,
    c::settings::OperationMode TOperationMode,
    c::settings::IntegrationMode TIntegrationMode
>
absl::Status PreprocessingForegroundContainer<TDeviceType,
    TOperationMode,
    TIntegrationMode>::HandleOutputData(int64_t frame_timestamp) {
    bool got_metrics_output;
    nlohmann::json preprocessed_data;
    MP_RETURN_IF_ERROR(ph::GetPacketContentsIfAny(
        preprocessed_data, got_metrics_output, preprocessed_data_poller.Get(), pg::output_streams::kPreprocessedData,
        this->settings.verbosity_level > 1
    ));
    if (got_metrics_output) {
        MP_RETURN_IF_ERROR(this->OnPreprocessedDataOutput(preprocessed_data, frame_timestamp));
        if (TOperationMode == c::settings::OperationMode::Spot) {
            // reset to start state
            this->recording = false;
            if (this->load_video) {
                this->keep_grabbing_frames = false;
            } else if (this->settings.video_source.auto_lock &&
                       this->video_source->SupportsExposureControls()) {
                MP_RETURN_IF_ERROR(this->video_source->TurnOnAutoExposure());
            }
        }
    }
    return absl::OkStatus();
}
 
template<
    platform_independence::DeviceType TDeviceType,
    container::settings::OperationMode TOperationMode,
    container::settings::IntegrationMode TIntegrationMode
>
std::string PreprocessingForegroundContainer<TDeviceType,TOperationMode,TIntegrationMode>::GetThirdGraphFileSuffix()
const {
    return c::settings::IntegrationModeTraits<TIntegrationMode>::kPreprocessingDataFormat;
}
 
template<
    platform_independence::DeviceType TDeviceType,
    container::settings::OperationMode TOperationMode,
    container::settings::IntegrationMode TIntegrationMode
>
std::string PreprocessingForegroundContainer<TDeviceType,TOperationMode,TIntegrationMode>::GetGraphFilePrefix() const {
    return "preprocessing";
}
 
template<
    platform_independence::DeviceType TDeviceType,
    c::settings::OperationMode TOperationMode,
    c::settings::IntegrationMode TIntegrationMode
>
absl::Status PreprocessingForegroundContainer<TDeviceType,
    TOperationMode,
    TIntegrationMode>::DrawDebugOverlays(cv::Mat& output_frame_bgr) {
    using ::presage::physiology::Landmarks;
 
    bool got_packet;
 
    // Verbosity and transformed size capture (once)
    const bool verbose = this->Base::settings.verbosity_level > 3;
    if (!this->transformed_size.has_value()) {
        std::pair<int, int> ts;
        bool size_available;
        MP_RETURN_IF_ERROR(ph::GetPacketContentsIfAny(
            ts, size_available, this->transformed_size_poller.Get(), pg::output_streams::kTransformedSize, verbose
        ));
        if (size_available) {
            this->transformed_size = cv::Size(ts.first, ts.second);
        }
    }
    // Precompute scale factors (multiply inside point loop only)
    double scale_x = 1.0, scale_y = 1.0;
    if (this->transformed_size.has_value() && output_frame_bgr.data &&
        this->transformed_size->width > 0 && this->transformed_size->height > 0) {
        scale_x = static_cast<double>(output_frame_bgr.cols) / static_cast<double>(this->transformed_size->width);
        scale_y = static_cast<double>(output_frame_bgr.rows) / static_cast<double>(this->transformed_size->height);
    }
 
    auto draw_region = [&, this](const Landmarks& landmarks, const cv::Scalar& color, const std::string& label, double alpha) {
        if (!output_frame_bgr.data || landmarks.value_size() == 0) return;
        std::vector<cv::Point> points;
        points.reserve(landmarks.value_size());
        for (int i = 0; i < landmarks.value_size(); ++i) {
            // 0.0f is magic value for "undefined"
            if (landmarks.value(i).x() != 0.0f){
                int x = cvRound(landmarks.value(i).x() * scale_x);
                int y = cvRound(landmarks.value(i).y() * scale_y);
                points.emplace_back(cv::Point(x, y));
            }
        }
        for (const auto& point : points) {
            cv::circle(output_frame_bgr, point, 3, color, cv::FILLED);
        }
        if (points.size() >= 3) {
            std::vector<cv::Point> hull;
            cv::convexHull(points, hull); // compute hull points directly
            cv::Mat mask = cv::Mat::zeros(output_frame_bgr.size(), CV_8UC1);
            std::vector<std::vector<cv::Point>> hulls = {hull};
            cv::fillPoly(mask, hulls, cv::Scalar(255));
            cv::Mat color_mask;
            cv::cvtColor(mask, color_mask, cv::COLOR_GRAY2BGR);
            color_mask.setTo(color, mask);
            cv::addWeighted(color_mask, alpha, output_frame_bgr, 1.0 - alpha, 0, output_frame_bgr);
            // Outline for clarity
            const int thickness = 2;
            cv::polylines(output_frame_bgr, hulls, true, color, thickness);
            // Label near centroid
            cv::Point2d centroid(0, 0);
            // conditional here on the off-chance our hull computation errored somehow and got 0 points
            if (hull.size() > 0){
                for(const auto& point : hull){
                    centroid.x += point.x;
                    centroid.y += point.y;
                }
                centroid.x /= static_cast<double>(hull.size());
                centroid.y /= static_cast<double>(hull.size());
            }
            cv::putText(output_frame_bgr, label, cv::Point(static_cast<int>(centroid.x), static_cast<int>(centroid.y)),
                        cv::FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv::LINE_AA);
        }
    };
 
    auto draw_if_available = [&](auto& poller,
                                 const char* stream_name,
                                 const char* label,
                                 const cv::Scalar& color,
                                 double alpha) -> absl::Status {
        Landmarks region_points;
        bool available;
        MP_RETURN_IF_ERROR(ph::GetPacketContentsIfAny(
            region_points, available, poller.Get(), stream_name, verbose
        ));
        if (available) draw_region(region_points, color, label, alpha);
        return absl::OkStatus();
    };
 
    MP_RETURN_IF_ERROR(draw_if_available(this->points_chest_poller,   pg::output_streams::kPointsChest,   "chest",   cv::Scalar(0, 0, 255),   0.30));
    MP_RETURN_IF_ERROR(draw_if_available(this->points_abdomen_poller, pg::output_streams::kPointsAbdomen, "abdomen", cv::Scalar(0, 255, 0),   0.25));
    MP_RETURN_IF_ERROR(draw_if_available(this->points_glutes_poller,  pg::output_streams::kPointsGlutes,  "glutes",  cv::Scalar(255, 0, 0),  0.20));
    MP_RETURN_IF_ERROR(draw_if_available(this->points_knees_poller,   pg::output_streams::kPointsKnees,   "knees",   cv::Scalar(0, 255, 255), 0.20));
 
    return absl::OkStatus();
}
 
} // namespace presage::smartspectra::lab