chicken-counting-sukawarna-det/cpp/include/chicken_counter/counting.hpp

#pragma once

#include <cstdint>
#include <deque>
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <opencv2/imgproc.hpp>
#include <opencv2/core/types.hpp>

#include "chicken_counter/config.hpp"
#include "chicken_counter/types.hpp"

namespace cc {

class CountingZone {
public:
    CountingZone() {}

    CountingZone(const RoiConfig& roi,
                 const GateConfig& gate,
                 int trail_length,
                 int track_buffer,
                 int min_box_area_px = 0,
                 bool validate_while_inside = true,
                 bool verbose = false)
        : roi(roi), gate(gate), trail_length(trail_length),
          track_buffer(track_buffer), min_box_area_px(min_box_area_px),
          min_overlap_ratio(roi.min_overlap_ratio),
          validate_while_inside(validate_while_inside),
          verbose(verbose)
    {
        for (auto& p : roi.counting_polygon())
            counting_polygon.push_back(p);
        auto r = roi.counting_rect();
        counting_rect = r;
    }

    std::vector<CountEvent> update(
        const std::vector<TrackObservation>& tracks,
        int frame_index,
        bool counting_paused = false)
    {
        std::vector<CountEvent> events;
        std::unordered_set<int> active_ids, inside_ids;

        for (const auto& track : tracks) {
            active_ids.insert(track.track_id);
            last_seen_frame[track.track_id] = frame_index;
            auto& hist = histories[track.track_id];
            if (hist.size() >= static_cast<size_t>(trail_length))
                hist.pop_front();
            hist.push_back({track.centroid_x, track.centroid_y});

            if (inside_roi({track.centroid_x, track.centroid_y}))
                inside_ids.insert(track.track_id);

            if (counting_paused) continue;
            if (inside_ids.find(track.track_id) == inside_ids.end()) continue;
            if (counted_ids.find(track.track_id) != counted_ids.end()) continue;

            bool should_validate = false;
            if (validate_while_inside) {
                should_validate = meets_validation_thresholds(track);
            } else {
                bool just_entered = inside_ids.find(track.track_id) != inside_ids.end()
                    && prev_inside_ids.find(track.track_id) == prev_inside_ids.end();
                should_validate = just_entered && meets_validation_thresholds(track);
            }

            if (should_validate) {
                counted_ids.insert(track.track_id);
                ++total_entered_count;
                sequence_numbers[track.track_id] = total_entered_count;
                latest_validated_track_id = track.track_id;
                CountEvent ev;
                ev.track_id = track.track_id;
                ev.frame_index = frame_index;
                ev.total_entered_after_event = total_entered_count;
                ev.sequence_number = total_entered_count;
                events.push_back(ev);

                if (verbose) {
                    int bbox_area = std::max(0, track.bbox_x2 - track.bbox_x1)
                                  * std::max(0, track.bbox_y2 - track.bbox_y1);
                    double overlap = bbox_overlap_ratio(track);
                    std::cerr << "[count] track=" << track.track_id
                              << " seq=#" << total_entered_count
                              << " frame=" << frame_index
                              << " area=" << bbox_area
                              << " overlap=" << overlap
                              << " conf=" << track.confidence
                              << " centroid=" << track.centroid_x << "," << track.centroid_y << "\n";
                }
            }
        }

        inside_box_count = static_cast<int>(inside_ids.size());
        current_inside_ids = inside_ids;
        prev_inside_ids = inside_ids;
        purge_stale(frame_index, active_ids);
        return events;
    }

    std::vector<cv::Point2i> trail_for(int track_id) const {
        auto it = histories.find(track_id);
        if (it == histories.end()) return {};
        return {it->second.begin(), it->second.end()};
    }

    size_t sequence_number_for(int track_id) const {
        auto it = sequence_numbers.find(track_id);
        return (it != sequence_numbers.end()) ? it->second : 0;
    }

    bool is_inside(int track_id) const {
        return current_inside_ids.find(track_id) != current_inside_ids.end();
    }

    bool is_validated(int track_id) const {
        return counted_ids.find(track_id) != counted_ids.end();
    }

    int inside_box_count = 0;
    int total_entered_count = 0;
    std::optional<int> latest_validated_track_id;

private:
    RoiConfig roi;
    GateConfig gate;
    int trail_length;
    int track_buffer;
    int min_box_area_px;
    float min_overlap_ratio;
    bool validate_while_inside;
    bool verbose;

    std::vector<cv::Point2i> counting_polygon;
    cv::Rect counting_rect;
    std::unordered_map<int, std::deque<cv::Point2i>> histories;
    std::unordered_map<int, int> last_seen_frame;
    std::unordered_set<int> counted_ids;
    std::unordered_set<int> prev_inside_ids;
    std::unordered_set<int> current_inside_ids;
    std::unordered_map<int, int> sequence_numbers;

    bool inside_roi(cv::Point2i p) const {
        return cv::pointPolygonTest(counting_polygon, p, false) > 0;
    }

    bool meets_size_threshold(const TrackObservation& track) const {
        int area = std::max(0, track.bbox_x2 - track.bbox_x1)
                 * std::max(0, track.bbox_y2 - track.bbox_y1);
        return area >= min_box_area_px;
    }

    double bbox_overlap_ratio(const TrackObservation& track) const {
        int bbox_area = std::max(0, track.bbox_x2 - track.bbox_x1)
                      * std::max(0, track.bbox_y2 - track.bbox_y1);
        if (bbox_area <= 0) return 0.0;

        int ix1 = std::max(track.bbox_x1, counting_rect.x);
        int iy1 = std::max(track.bbox_y1, counting_rect.y);
        int ix2 = std::min(track.bbox_x2, counting_rect.x + counting_rect.width);
        int iy2 = std::min(track.bbox_y2, counting_rect.y + counting_rect.height);
        if (ix2 <= ix1 || iy2 <= iy1) return 0.0;

        double intersection = (ix2 - ix1) * (iy2 - iy1);
        return intersection / bbox_area;
    }

    bool meets_overlap_threshold(const TrackObservation& track) const {
        if (min_overlap_ratio <= 0) return true;
        return bbox_overlap_ratio(track) >= min_overlap_ratio;
    }

    bool meets_validation_thresholds(const TrackObservation& track) const {
        return meets_size_threshold(track) && meets_overlap_threshold(track);
    }

    void purge_stale(int frame_index, const std::unordered_set<int>& active_ids) {
        std::vector<int> stale;
        for (const auto& [id, last] : last_seen_frame) {
            if (active_ids.find(id) == active_ids.end()
                && frame_index - last > track_buffer)
                stale.push_back(id);
        }
        for (int id : stale) {
            last_seen_frame.erase(id);
            histories.erase(id);
            prev_inside_ids.erase(id);
            current_inside_ids.erase(id);
        }
    }
};

} // namespace cc