forked from dsutanto/bytetrack-counter-cpp
438 lines
18 KiB
C++
438 lines
18 KiB
C++
#include "config.hpp"
|
|
#include "bytetrack.hpp"
|
|
#include "batch_store.hpp"
|
|
#include "csv_logger.hpp"
|
|
#include "rknn_yolo.hpp"
|
|
#include "drawing.hpp"
|
|
#include "nms.hpp"
|
|
|
|
#include <opencv2/opencv.hpp>
|
|
#include <Eigen/Dense>
|
|
#include <csignal>
|
|
#include <iostream>
|
|
#include <chrono>
|
|
#include <thread>
|
|
#include <unordered_map>
|
|
#include <unordered_set>
|
|
#include <deque>
|
|
#include <ctime>
|
|
#include <filesystem>
|
|
|
|
static std::atomic<bool> shutdown_requested{false};
|
|
|
|
static void signal_handler(int) {
|
|
shutdown_requested = true;
|
|
std::cout << "\nShutdown requested — finishing current frame..." << std::endl;
|
|
}
|
|
|
|
static std::string now_str() {
|
|
auto now = std::chrono::system_clock::now();
|
|
auto t = std::chrono::system_clock::to_time_t(now);
|
|
std::ostringstream ss;
|
|
ss << std::put_time(std::localtime(&t), "%Y-%m-%d %H:%M:%S");
|
|
return ss.str();
|
|
}
|
|
|
|
static cv::VideoCapture open_capture(const std::string& source) {
|
|
bool is_live = source.substr(0, 7) == "rtsp://" || source.substr(0, 7) == "http://";
|
|
if (is_live) {
|
|
setenv("OPENCV_FFMPEG_CAPTURE_OPTIONS", g_config.rtsp_ffmpeg_options.c_str(), 1);
|
|
}
|
|
cv::VideoCapture cap(source, cv::CAP_FFMPEG);
|
|
cap.set(cv::CAP_PROP_BUFFERSIZE, 1);
|
|
return cap;
|
|
}
|
|
|
|
static void warmup_stream(cv::VideoCapture& cap, int n) {
|
|
std::cout << "Warming up stream..." << std::endl;
|
|
cv::Mat frame;
|
|
for (int i = 0; i < n; ++i) cap.read(frame);
|
|
std::cout << "Stream ready!" << std::endl;
|
|
}
|
|
|
|
struct TrackedInfo {
|
|
float cx;
|
|
std::chrono::steady_clock::time_point ts;
|
|
};
|
|
|
|
static void prune_stale_tracks(std::unordered_map<int, TrackedInfo>& tracked, double prune_sec) {
|
|
auto now = std::chrono::steady_clock::now();
|
|
std::vector<int> stale;
|
|
for (auto& [tid, info] : tracked) {
|
|
auto age = std::chrono::duration_cast<std::chrono::seconds>(now - info.ts).count();
|
|
if (age > prune_sec) stale.push_back(tid);
|
|
}
|
|
for (int tid : stale) tracked.erase(tid);
|
|
}
|
|
|
|
int main(int argc, char* argv[]) {
|
|
signal(SIGINT, signal_handler);
|
|
signal(SIGTERM, signal_handler);
|
|
|
|
const char* env_path = argc > 1 ? argv[1] : nullptr;
|
|
g_config.load_from_env(env_path);
|
|
|
|
std::cout << "Starting ByteTrack Counter (C++)" << std::endl;
|
|
std::cout << "Camera: " << g_config.camera_name << std::endl;
|
|
std::cout << "Source: " << g_config.source << std::endl;
|
|
std::cout << "Model: " << g_config.model_path << std::endl;
|
|
std::cout << "DB: " << g_config.db_path << std::endl;
|
|
|
|
int class_ayam = 0;
|
|
int class_talenan = 1;
|
|
|
|
auto logger = [](const std::string& msg) {
|
|
std::cout << "[" << now_str() << "] " << msg << std::endl;
|
|
};
|
|
|
|
BatchStore store(g_config.db_path, g_config.state_file,
|
|
g_config.camera_name, g_config.object_label,
|
|
g_config.daily_cutoff_time,
|
|
g_config.batch_timeout_seconds,
|
|
g_config.ignore_batch_label_timeout,
|
|
g_config.min_object_per_batch,
|
|
g_config.min_duration_per_batch,
|
|
logger);
|
|
store.start_cutoff_watcher();
|
|
|
|
std::unique_ptr<CsvLogger> cross_logger;
|
|
if (g_config.export_csv) {
|
|
cross_logger = std::make_unique<CsvLogger>(g_config.cross_csv,
|
|
std::vector<std::string>{"batch", "frame", "timestamp", "chicken_id"});
|
|
}
|
|
|
|
RKNNYOLO model(g_config.model_path, g_config.core_mask,
|
|
g_config.imgsz, g_config.conf, 0.45f,
|
|
g_config.num_classes, 0, g_config.score_sigmoid);
|
|
|
|
ByteTracker ayam_tracker(g_config.track_high_thresh, g_config.track_low_thresh,
|
|
g_config.track_match_thresh, g_config.track_buffer,
|
|
g_config.track_min_hits);
|
|
ByteTracker talenan_tracker(g_config.track_high_thresh, g_config.track_low_thresh,
|
|
g_config.track_match_thresh, g_config.track_buffer,
|
|
g_config.track_min_hits);
|
|
|
|
std::unordered_set<int> ayam_line_crossed;
|
|
std::unordered_set<int> talenan_line_crossed;
|
|
|
|
std::unordered_map<int, int> ayam_cross_flash;
|
|
std::unordered_map<int, int> talenan_cross_flash;
|
|
int line_pulse = 0, count_pulse = 0, batch_pulse = 0;
|
|
std::vector<Popup> popups;
|
|
|
|
auto session_start = std::chrono::steady_clock::now();
|
|
int frame_idx = 0;
|
|
|
|
cv::VideoCapture cap = open_capture(g_config.source);
|
|
bool is_live = g_config.source.substr(0, 7) == "rtsp://" ||
|
|
g_config.source.substr(0, 7) == "http://";
|
|
|
|
if (!cap.isOpened()) {
|
|
std::cerr << "Cannot open source: " << g_config.source << std::endl;
|
|
store.shutdown();
|
|
return 1;
|
|
}
|
|
if (is_live && g_config.warmup_frames > 0)
|
|
warmup_stream(cap, g_config.warmup_frames);
|
|
|
|
int w = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
|
|
int h = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
|
|
double fps_prop = cap.get(cv::CAP_PROP_FPS);
|
|
if (fps_prop <= 1) fps_prop = g_config.output_fps;
|
|
|
|
int line_x = resolve_line_x(w, g_config.line_x, g_config.line_x_frac);
|
|
|
|
std::cout << "RKNN+ByteTrack counter | " << w << "x" << h << " @ "
|
|
<< fps_prop << "fps | line x=" << line_x
|
|
<< " | cross=" << g_config.cross_direction << std::endl;
|
|
|
|
int reconnect_count = 0;
|
|
|
|
std::unordered_map<int, TrackedInfo> ayam_tracked;
|
|
std::unordered_map<int, TrackedInfo> talenan_tracked;
|
|
|
|
while (!shutdown_requested) {
|
|
cv::Mat frame;
|
|
cap >> frame;
|
|
if (frame.empty()) {
|
|
if (!is_live) break;
|
|
reconnect_count++;
|
|
std::cout << "Stream dropped (attempt " << reconnect_count
|
|
<< "), reconnecting..." << std::endl;
|
|
cap.release();
|
|
std::this_thread::sleep_for(std::chrono::seconds(g_config.reconnect_delay_sec));
|
|
cap = open_capture(g_config.source);
|
|
if (!cap.isOpened()) break;
|
|
w = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
|
|
h = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
|
|
line_x = resolve_line_x(w, g_config.line_x, g_config.line_x_frac);
|
|
continue;
|
|
}
|
|
|
|
auto elapsed_sec = std::chrono::duration_cast<std::chrono::seconds>(
|
|
std::chrono::steady_clock::now() - session_start).count();
|
|
auto mono = std::chrono::steady_clock::now();
|
|
|
|
bool ayam_crossed_frame = false, batch_closed_frame = false, batch_started_frame = false;
|
|
|
|
std::vector<Detection> detections = model(frame);
|
|
|
|
if (!detections.empty()) {
|
|
std::vector<Eigen::Vector4f> ayam_boxes, talenan_boxes;
|
|
std::vector<float> ayam_scores, talenan_scores;
|
|
std::vector<float> ayam_cx_list, talenan_cx_list;
|
|
std::vector<std::vector<std::vector<cv::Point2f>>> ayam_kpts_list;
|
|
|
|
for (const auto& det : detections) {
|
|
const auto& bbox = det.bbox;
|
|
float cx = (bbox(0) + bbox(2)) / 2.0f;
|
|
|
|
if (det.cls == class_talenan) {
|
|
talenan_boxes.push_back(bbox);
|
|
talenan_scores.push_back(det.score);
|
|
talenan_cx_list.push_back(cx);
|
|
} else if (det.cls == class_ayam) {
|
|
ayam_boxes.push_back(bbox);
|
|
ayam_scores.push_back(det.score);
|
|
ayam_cx_list.push_back(cx);
|
|
ayam_kpts_list.push_back(det.keypoints);
|
|
}
|
|
}
|
|
|
|
int na = static_cast<int>(ayam_boxes.size());
|
|
int nt = static_cast<int>(talenan_boxes.size());
|
|
|
|
Eigen::MatrixXf ayam_mat(na, 4), talenan_mat(nt, 4);
|
|
Eigen::VectorXf ayam_scores_vec(na), talenan_scores_vec(nt);
|
|
|
|
for (int i = 0; i < na; ++i) {
|
|
ayam_mat.row(i) = ayam_boxes[i];
|
|
ayam_scores_vec(i) = ayam_scores[i];
|
|
}
|
|
for (int i = 0; i < nt; ++i) {
|
|
talenan_mat.row(i) = talenan_boxes[i];
|
|
talenan_scores_vec(i) = talenan_scores[i];
|
|
}
|
|
|
|
// Apply NMS per-class to suppress duplicate boxes
|
|
if (na > 0) {
|
|
auto keep = nms(ayam_mat, ayam_scores_vec);
|
|
int nk = static_cast<int>(keep.size());
|
|
Eigen::MatrixXf fm(nk, 4);
|
|
Eigen::VectorXf fs(nk);
|
|
std::vector<Eigen::Vector4f> fb;
|
|
std::vector<float> fcx;
|
|
std::vector<std::vector<std::vector<cv::Point2f>>> fkpts;
|
|
for (int ki = 0; ki < nk; ++ki) {
|
|
int idx = keep[ki];
|
|
fm.row(ki) = ayam_mat.row(idx);
|
|
fs(ki) = ayam_scores_vec(idx);
|
|
fb.push_back(ayam_boxes[idx]);
|
|
fcx.push_back(ayam_cx_list[idx]);
|
|
if (idx < (int)ayam_kpts_list.size())
|
|
fkpts.push_back(ayam_kpts_list[idx]);
|
|
}
|
|
na = nk;
|
|
ayam_mat = std::move(fm);
|
|
ayam_scores_vec = std::move(fs);
|
|
ayam_boxes = std::move(fb);
|
|
ayam_cx_list = std::move(fcx);
|
|
ayam_kpts_list = std::move(fkpts);
|
|
}
|
|
if (nt > 0) {
|
|
auto keep = nms(talenan_mat, talenan_scores_vec);
|
|
int nk = static_cast<int>(keep.size());
|
|
Eigen::MatrixXf fm(nk, 4);
|
|
Eigen::VectorXf fs(nk);
|
|
std::vector<Eigen::Vector4f> fb;
|
|
std::vector<float> fcx;
|
|
for (int ki = 0; ki < nk; ++ki) {
|
|
int idx = keep[ki];
|
|
fm.row(ki) = talenan_mat.row(idx);
|
|
fs(ki) = talenan_scores_vec(idx);
|
|
fb.push_back(talenan_boxes[idx]);
|
|
fcx.push_back(talenan_cx_list[idx]);
|
|
}
|
|
nt = nk;
|
|
talenan_mat = std::move(fm);
|
|
talenan_scores_vec = std::move(fs);
|
|
talenan_boxes = std::move(fb);
|
|
talenan_cx_list = std::move(fcx);
|
|
}
|
|
|
|
TrackResult ayam_tr = ayam_tracker.update(ayam_mat, ayam_scores_vec);
|
|
TrackResult talenan_tr = talenan_tracker.update(talenan_mat, talenan_scores_vec);
|
|
|
|
// Process talenan crossings
|
|
for (int di = 0; di < nt; ++di) {
|
|
auto it = talenan_tr.det_to_track.find(di);
|
|
if (it == talenan_tr.det_to_track.end()) continue;
|
|
int tid = it->second;
|
|
float cx = talenan_cx_list[di];
|
|
|
|
auto tit = talenan_tracked.find(tid);
|
|
if (tit != talenan_tracked.end()) {
|
|
float prev_cx = tit->second.cx;
|
|
if (crossed_line(prev_cx, cx, line_x, g_config.cross_direction) &&
|
|
talenan_line_crossed.find(tid) == talenan_line_crossed.end()) {
|
|
talenan_line_crossed.insert(tid);
|
|
if (store.record_talenan_crossing(tid)) {
|
|
batch_closed_frame = true;
|
|
talenan_cross_flash[tid] = g_config.cross_flash_frames;
|
|
popups.push_back({
|
|
static_cast<int>(cx) - 20,
|
|
static_cast<int>((talenan_boxes[di](1) + talenan_boxes[di](3)) / 2),
|
|
frame_idx, "BATCH CLOSED"
|
|
});
|
|
}
|
|
}
|
|
}
|
|
talenan_tracked[tid] = {cx, mono};
|
|
}
|
|
|
|
// Process ayam crossings
|
|
for (int di = 0; di < na; ++di) {
|
|
auto it = ayam_tr.det_to_track.find(di);
|
|
if (it == ayam_tr.det_to_track.end()) continue;
|
|
int tid = it->second;
|
|
float cx = ayam_cx_list[di];
|
|
|
|
auto tit = ayam_tracked.find(tid);
|
|
if (tit != ayam_tracked.end()) {
|
|
float prev_cx = tit->second.cx;
|
|
if (crossed_line(prev_cx, cx, line_x, g_config.cross_direction) &&
|
|
ayam_line_crossed.find(tid) == ayam_line_crossed.end()) {
|
|
ayam_line_crossed.insert(tid);
|
|
auto [count, started_new] = store.record_ayam_crossing(tid);
|
|
if (cross_logger) {
|
|
std::ostringstream frame_str, tid_str;
|
|
frame_str << frame_idx;
|
|
tid_str << tid;
|
|
cross_logger->write_row({
|
|
std::to_string(store.current_batch_number()),
|
|
frame_str.str(),
|
|
now_str(),
|
|
tid_str.str()
|
|
});
|
|
}
|
|
ayam_crossed_frame = true;
|
|
if (started_new) batch_started_frame = true;
|
|
ayam_cross_flash[tid] = g_config.cross_flash_frames;
|
|
popups.push_back({
|
|
static_cast<int>(cx) - 12,
|
|
static_cast<int>((ayam_boxes[di](1) + ayam_boxes[di](3)) / 2),
|
|
frame_idx, "+1"
|
|
});
|
|
}
|
|
}
|
|
ayam_tracked[tid] = {cx, mono};
|
|
}
|
|
|
|
// Track lost tracks for continuity
|
|
for (auto& [tid, cx] : ayam_tr.lost_map) {
|
|
if (ayam_tracked.find(tid) == ayam_tracked.end())
|
|
ayam_tracked[tid] = {cx, mono};
|
|
}
|
|
|
|
// Draw talenan
|
|
for (int di = 0; di < nt; ++di) {
|
|
auto it = talenan_tr.det_to_track.find(di);
|
|
if (it == talenan_tr.det_to_track.end()) continue;
|
|
int tid = it->second;
|
|
const auto& bbox = talenan_boxes[di];
|
|
int x1 = static_cast<int>(bbox(0)), y1 = static_cast<int>(bbox(1));
|
|
int x2 = static_cast<int>(bbox(2)), y2 = static_cast<int>(bbox(3));
|
|
int flash = talenan_cross_flash.count(tid) ? talenan_cross_flash[tid] : 0;
|
|
cv::Scalar color = flash > 0 ? C_GREEN : C_TALENAN_BOX;
|
|
cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, flash > 0 ? 3 : 2);
|
|
draw_pill(frame, "TALENAN " + std::to_string(tid), x1, y1 - 4, color);
|
|
}
|
|
|
|
// Draw ayam
|
|
for (int di = 0; di < na; ++di) {
|
|
auto it = ayam_tr.det_to_track.find(di);
|
|
if (it == ayam_tr.det_to_track.end()) continue;
|
|
int tid = it->second;
|
|
const auto& bbox = ayam_boxes[di];
|
|
int x1 = static_cast<int>(bbox(0)), y1 = static_cast<int>(bbox(1));
|
|
int x2 = static_cast<int>(bbox(2)), y2 = static_cast<int>(bbox(3));
|
|
int flash = ayam_cross_flash.count(tid) ? ayam_cross_flash[tid] : 0;
|
|
cv::Scalar color = flash > 0 ? C_GREEN : C_AYAM_BOX;
|
|
cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, flash > 0 ? 3 : 2);
|
|
draw_pill(frame, "ID " + std::to_string(tid), x1, y1 - 4, color);
|
|
if (di < (int)ayam_kpts_list.size() && !ayam_kpts_list[di].empty())
|
|
draw_skeleton_bold(frame, ayam_kpts_list[di]);
|
|
}
|
|
}
|
|
|
|
if (ayam_crossed_frame) {
|
|
line_pulse = g_config.line_pulse_frames;
|
|
count_pulse = g_config.count_pulse_frames;
|
|
}
|
|
if (batch_closed_frame) {
|
|
line_pulse = g_config.line_pulse_frames;
|
|
}
|
|
if (batch_started_frame) {
|
|
batch_pulse = g_config.batch_pulse_frames;
|
|
}
|
|
|
|
int batch_num = store.current_batch_number();
|
|
int batch_count = store.current_batch_count();
|
|
int display_total = store.display_total();
|
|
double elapsed = static_cast<double>(elapsed_sec);
|
|
double rate = elapsed > 0 ? (display_total / elapsed * 60.0) : 0.0;
|
|
|
|
draw_elegant_counting_line(frame, line_x, h, line_pulse);
|
|
draw_hero_count(frame, line_x, h, batch_count, count_pulse);
|
|
draw_hud(frame, w, batch_num, batch_count, display_total,
|
|
elapsed, rate, g_config.camera_name, now_str());
|
|
draw_batch_banner(frame, w, batch_num, batch_pulse);
|
|
draw_footer(frame, w, h, frame_idx, is_live ? "LIVE-RKNN-BT" : "FILE-RKNN-BT");
|
|
popups = draw_popups(frame, popups, frame_idx);
|
|
|
|
// Flash decay
|
|
for (auto& flash_map : {&ayam_cross_flash, &talenan_cross_flash}) {
|
|
for (auto it = flash_map->begin(); it != flash_map->end(); ) {
|
|
it->second--;
|
|
if (it->second <= 0) it = flash_map->erase(it);
|
|
else ++it;
|
|
}
|
|
}
|
|
line_pulse = std::max(0, line_pulse - 1);
|
|
count_pulse = std::max(0, count_pulse - 1);
|
|
batch_pulse = std::max(0, batch_pulse - 1);
|
|
|
|
if (g_config.live_stream_enabled && frame_idx % g_config.live_stream_every_n == 0) {
|
|
try {
|
|
std::filesystem::path p(g_config.live_stream_frame_path);
|
|
if (p.has_parent_path())
|
|
std::filesystem::create_directories(p.parent_path());
|
|
std::vector<uint8_t> jpeg;
|
|
cv::imencode(".jpg", frame, jpeg,
|
|
{cv::IMWRITE_JPEG_QUALITY, g_config.live_stream_quality});
|
|
std::ofstream f(g_config.live_stream_frame_path, std::ios::binary);
|
|
f.write(reinterpret_cast<const char*>(jpeg.data()), jpeg.size());
|
|
} catch (...) {}
|
|
}
|
|
|
|
frame_idx++;
|
|
if (frame_idx % g_config.flush_every_n_frames == 0) {
|
|
std::cout << "[" << now_str() << "] Frame " << frame_idx
|
|
<< " | Batch " << batch_num << ": " << batch_count
|
|
<< " | Total: " << display_total
|
|
<< " | Uptime " << (elapsed / 3600.0) << "h" << std::endl;
|
|
}
|
|
|
|
prune_stale_tracks(ayam_tracked, g_config.tracked_prune_sec);
|
|
prune_stale_tracks(talenan_tracked, g_config.tracked_prune_sec);
|
|
}
|
|
|
|
cap.release();
|
|
store.shutdown();
|
|
|
|
std::cout << "\n=== Batch Summary (SQLite) ===" << std::endl;
|
|
std::cout << "Database: " << g_config.db_path << std::endl;
|
|
|
|
return 0;
|
|
}
|