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

#pragma once

#include <cstdint>
#include <string>

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

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

namespace cc {

const cv::Scalar WHITE(255, 255, 255);
const cv::Scalar BLACK(0, 0, 0);
const cv::Scalar CYAN(255, 255, 0);
const cv::Scalar RED(0, 0, 255);
const cv::Scalar YELLOW(0, 255, 255);
const cv::Scalar ORANGE(0, 165, 255);
const cv::Scalar BLUE(255, 120, 0);
const cv::Scalar LIME(80, 220, 80);
const cv::Scalar GRAY(160, 160, 160);

inline void draw_outlined_text(
    cv::Mat& frame,
    const std::string& text,
    cv::Point origin,
    double font_scale,
    const cv::Scalar& fill_color,
    const cv::Scalar& outline_color,
    int thickness,
    int outline_thickness)
{
    int font = cv::FONT_HERSHEY_SIMPLEX;
    cv::putText(frame, text, origin, font, font_scale, outline_color, outline_thickness, cv::LINE_AA);
    cv::putText(frame, text, origin, font, font_scale, fill_color, thickness, cv::LINE_AA);
}

inline void draw_trail(cv::Mat& frame, const std::vector<cv::Point2i>& trail) {
    if (trail.size() < 2) return;
    for (size_t i = 1; i < trail.size(); ++i)
        cv::line(frame, trail[i - 1], trail[i], YELLOW, 2);
}

inline void draw_dashed_rectangle(
    cv::Mat& frame,
    cv::Point pt1,
    cv::Point pt2,
    const cv::Scalar& color,
    int thickness = 1,
    int dash_length = 12)
{
    int x1 = pt1.x, y1 = pt1.y, x2 = pt2.x, y2 = pt2.y;
    for (int x = x1; x < x2; x += dash_length * 2)
        cv::line(frame, {x, y1}, {std::min(x + dash_length, x2), y1}, color, thickness);
    for (int x = x1; x < x2; x += dash_length * 2)
        cv::line(frame, {x, y2}, {std::min(x + dash_length, x2), y2}, color, thickness);
    for (int y = y1; y < y2; y += dash_length * 2)
        cv::line(frame, {x1, y}, {x1, std::min(y + dash_length, y2)}, color, thickness);
    for (int y = y1; y < y2; y += dash_length * 2)
        cv::line(frame, {x2, y}, {x2, std::min(y + dash_length, y2)}, color, thickness);
}

inline void draw_detection_zone(cv::Mat& frame, const CameraConfig& config) {
    int h = frame.rows, w = frame.cols;
    auto r = config.detection_zone.compute_rect(config.roi, w, h);
    draw_dashed_rectangle(frame, {r.x, r.y}, {r.x + r.width, r.y + r.height}, GRAY, 1);
}

inline void draw_roi_and_gates(cv::Mat& frame, const CameraConfig& config) {
    std::vector<cv::Point2i> pts = config.roi.counting_polygon();
    std::vector<std::vector<cv::Point>> contours(1);
    for (const auto& p : pts) contours[0].emplace_back(p);
    cv::polylines(frame, contours, true, BLUE, 3);
}

inline cv::Mat draw_overlay(
    cv::Mat& frame,
    const CameraConfig& config,
    CountingZone& counting_zone,
    const std::vector<TrackObservation>& tracks,
    const MotionState& motion_state,
    int frame_index = 0,
    cv::Mat* buffer = nullptr)
{
    cv::Mat annotated;
    if (buffer) {
        frame.copyTo(*buffer);
        annotated = *buffer;
    } else {
        annotated = frame.clone();
    }

    if (config.detection_zone.enabled && config.detection_zone.show_in_overlay)
        draw_detection_zone(annotated, config);

    draw_roi_and_gates(annotated, config);

    bool blink_on = (frame_index / 8) % 2 == 0;
    const auto& pc = config.overlay.pending_colors;
    auto pending_colors = pc.empty()
        ? std::vector<cv::Scalar>{CYAN, YELLOW}
        : pc;

    for (const auto& track : tracks) {
        bool inside_box = counting_zone.is_inside(track.track_id);
        if (config.overlay.inside_box_only && !inside_box) continue;

        bool validated = counting_zone.is_validated(track.track_id);
        if (config.overlay.validated_only && !validated) continue;
        int x1 = track.bbox_x1, y1 = track.bbox_y1, x2 = track.bbox_x2, y2 = track.bbox_y2;
        int cx = track.centroid_x, cy = track.centroid_y;
        int seq = static_cast<int>(counting_zone.sequence_number_for(track.track_id));

        if (config.overlay.show_boxes) {
            cv::Scalar box_color;
            if (validated) {
                box_color = ORANGE;
            } else if (config.overlay.pending_blink) {
                box_color = pending_colors[blink_on ? 0 : 1 % pending_colors.size()];
            } else {
                box_color = pending_colors[0];
            }
            cv::rectangle(annotated, {x1, y1}, {x2, y2}, box_color, 2);
        }

        if (validated && seq > 0) {
            draw_outlined_text(annotated, std::to_string(seq),
                {x1, std::max(24, y1 - 8)}, 0.8, LIME, BLACK, 2, 4);
        }

        if (config.overlay.show_center_marker) {
            cv::Scalar marker_color = validated ? ORANGE
                : pending_colors[blink_on ? 0 : 1 % pending_colors.size()];
            cv::circle(annotated, {cx, cy}, 4, marker_color, -1);
            if (config.overlay.show_track_ring) {
                int radius = std::max(20, static_cast<int>(std::max(x2 - x1, y2 - y1) * 0.6));
                cv::circle(annotated, {cx, cy}, radius, WHITE, 1);
            }
        }

        if (config.overlay.show_track_trails) {
            auto trail = counting_zone.trail_for(track.track_id);
            draw_trail(annotated, trail);
        }
    }

    auto& anchor = config.overlay.count_anchor;
    draw_outlined_text(annotated,
        "TOTAL ENTERED: " + std::to_string(counting_zone.total_entered_count),
        anchor, 1.35, BLUE, BLACK, 4, 6);

    const char* motion_label = motion_state.backward_active ? "BACKWARD STOP" : "FORWARD";
    cv::Scalar motion_color = motion_state.backward_active ? RED : YELLOW;
    cv::putText(annotated, motion_label, {anchor.x, anchor.y + 42},
        cv::FONT_HERSHEY_SIMPLEX, 0.8, motion_color, 2, cv::LINE_AA);

    return annotated;
}

} // namespace cc