ariska/bytetrack-counter-cpp
1
0 Fork 0
forked from dsutanto/bytetrack-counter-cpp
Code Pull Requests Activity
Files
e67109ea88bd94e48b482a8f95e214ccf896655d
bytetrack-counter-cpp/src/main.cpp
T
proitlab e67109ea88 Recompile
2026-07-29 17:01:43 +07:00

602 lines
25 KiB
C++
Raw Blame History

#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 "version.hpp"
#include "video_writer.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;
}
static bool connect_stream(cv::VideoCapture& cap, int& w, int& h, double& fps) {
bool is_live = g_config.source.substr(0, 7) == "rtsp://" ||
g_config.source.substr(0, 7) == "http://";
int attempts = 0;
while (!shutdown_requested) {
cap = open_capture(g_config.source);
if (cap.isOpened()) break;
attempts++;
if (g_config.max_reconnect_attempts > 0 && attempts >= g_config.max_reconnect_attempts) {
std::cerr << "Cannot open source after " << attempts
<< " attempts: " << g_config.source << std::endl;
return false;
}
std::cout << "Cannot open source, retry in " << g_config.reconnect_delay_sec
<< "s..." << std::endl;
std::this_thread::sleep_for(std::chrono::seconds(g_config.reconnect_delay_sec));
}
if (shutdown_requested) return false;
if (is_live && g_config.warmup_frames > 0)
warmup_stream(cap, g_config.warmup_frames);
w = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
h = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
fps = cap.get(cv::CAP_PROP_FPS);
if (fps <= 1) fps = g_config.output_fps;
return true;
}
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);
if (argc < 2) {
std::cout << "ByteTrack Counter (C++)" << std::endl;
std::cout << "Version: " << version::GIT_COMMIT
<< " | branch: " << version::GIT_BRANCH
<< " | built: " << version::BUILD_DATE << " " << version::BUILD_TIME
<< " | rknn: " << (version::BUILD_RKNN ? "ON" : "OFF");
if (!version::GIT_TAG.empty())
std::cout << " | tag: " << version::GIT_TAG;
std::cout << std::endl;
std::cout << "Usage: bytetrack-counter <config.env>" << std::endl;
std::cout << " config.env path to environment configuration file (required)" << std::endl;
return 1;
}
g_config.load_from_env(argv[1]);
std::cout << "Starting ByteTrack Counter (C++)" << std::endl;
std::cout << "Version: " << version::GIT_COMMIT
<< " | branch: " << version::GIT_BRANCH
<< " | built: " << version::BUILD_DATE << " " << version::BUILD_TIME
<< " | rknn: " << (version::BUILD_RKNN ? "ON" : "OFF");
if (!version::GIT_TAG.empty())
std::cout << " | tag: " << version::GIT_TAG;
std::cout << 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 << " | imgsz=" << g_config.imgsz
<< " | core_mask=" << g_config.core_mask << std::endl;
std::cout << "DB: " << g_config.db_path << std::endl;
std::cout << "State: " << g_config.state_file << std::endl;
std::cout << "ByteTrack Index 0: high_thresh=" << g_config.track_high_thresh_0
<< " low_thresh=" << g_config.track_low_thresh_0
<< " match_thresh=" << g_config.track_match_thresh_0
<< " buffer=" << g_config.track_buffer_0
<< " min_hits=" << g_config.track_min_hits_0 << std::endl;
std::cout << "ByteTrack Index 1: high_thresh=" << g_config.track_high_thresh_1
<< " low_thresh=" << g_config.track_low_thresh_1
<< " match_thresh=" << g_config.track_match_thresh_1
<< " buffer=" << g_config.track_buffer_1
<< " min_hits=" << g_config.track_min_hits_1 << std::endl;
{
std::string reset_file = g_config.shm_dir + "/.reset";
if (std::filesystem::exists(reset_file)) {
std::cout << "[RESET] [INFO] .reset detected — clearing state" << std::endl;
std::error_code ec;
std::filesystem::remove(g_config.state_file, ec);
std::filesystem::remove(g_config.live_stream_frame_path, ec);
std::filesystem::remove(reset_file, ec);
}
}
if (g_config.live_stream_enabled) {
std::filesystem::path live_p(g_config.live_stream_frame_path);
if (live_p.has_parent_path())
std::filesystem::create_directories(live_p.parent_path());
}
int class_ayam = g_config.ayam_class_id;
int class_talenan = g_config.talenan_class_id;
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);
if (!model.is_loaded()) {
std::cerr << "ERROR: Failed to load model: " << g_config.model_path
<< std::endl;
return 1;
}
ByteTracker ayam_tracker(g_config.track_high_thresh_0, g_config.track_low_thresh_0,
g_config.track_match_thresh_0, g_config.track_buffer_0,
g_config.track_min_hits_0);
ByteTracker talenan_tracker(g_config.track_high_thresh_1, g_config.track_low_thresh_1,
g_config.track_match_thresh_1, g_config.track_buffer_1,
g_config.track_min_hits_1);
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;
std::deque<std::chrono::steady_clock::time_point> crossing_times;
auto session_start = std::chrono::steady_clock::now();
int frame_idx = 0;
float inf_ema = 0.0f;
std::string model_name = g_config.model_path;
auto slash = model_name.find_last_of("/\\");
if (slash != std::string::npos) model_name = model_name.substr(slash + 1);
cv::VideoCapture cap;
int w, h;
double fps_prop;
if (!connect_stream(cap, w, h, fps_prop)) {
store.shutdown();
return shutdown_requested ? 0 : 1;
}
bool is_live = g_config.source.substr(0, 7) == "rtsp://" ||
g_config.source.substr(0, 7) == "http://";
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;
std::unique_ptr<VideoBatchWriter> video_writer;
if (g_config.record_video) {
video_writer = std::make_unique<VideoBatchWriter>(
g_config.shm_dir, g_config.output_dir,
w, h, static_cast<int>(fps_prop),
g_config.record_delay_seconds,
g_config.record_end_delay_seconds,
g_config.record_retention_days);
auto* vw = video_writer.get();
store.set_batch_closed_callback([vw](int bn) { vw->on_batch_closed(bn); });
}
int reconnect_count = 0;
bool blank_live_frame = false;
std::unordered_map<int, TrackedInfo> ayam_tracked;
std::unordered_map<int, TrackedInfo> talenan_tracked;
cv::Mat prev_gray;
while (!shutdown_requested) {
{
std::string reset_file = g_config.shm_dir + "/.reset";
if (std::filesystem::exists(reset_file)) {
std::cout << "[" << now_str() << "] [INFO] [RESET] .reset detected — clearing state" << std::endl;
std::error_code ec;
std::filesystem::remove(g_config.state_file, ec);
std::filesystem::remove(g_config.live_stream_frame_path, ec);
std::filesystem::remove(reset_file, ec);
blank_live_frame = true;
}
}
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));
if (!connect_stream(cap, w, h, fps_prop)) break;
line_x = resolve_line_x(w, g_config.line_x, g_config.line_x_frac);
continue;
}
auto elapsed_sec = std::chrono::duration<double>(
std::chrono::steady_clock::now() - session_start).count();
auto mono = std::chrono::steady_clock::now();
if (store.check_cutoff_reset()) {
if (g_config.reset_counters_at_cutoff) {
std::cout << "[" << now_str() << "] Cutoff reset — resetting frame + tracker counters" << std::endl;
frame_idx = 0;
KalmanBoxTracker::count = 0;
ayam_tracker.reset();
talenan_tracker.reset();
ayam_tracked.clear();
talenan_tracked.clear();
ayam_line_crossed.clear();
talenan_line_crossed.clear();
ayam_cross_flash.clear();
talenan_cross_flash.clear();
crossing_times.clear();
line_pulse = 0;
count_pulse = 0;
batch_pulse = 0;
}
if (video_writer) video_writer->cleanup_retention();
}
bool ayam_crossed_frame = false, batch_closed_frame = false, batch_started_frame = false;
bool skip_inference = false;
if (g_config.motion_detection_enabled) {
cv::Mat gray;
cv::cvtColor(frame, gray, cv::COLOR_BGR2GRAY);
if (!prev_gray.empty()) {
cv::Mat diff;
cv::absdiff(gray, prev_gray, diff);
double mean_diff = cv::mean(diff)[0];
skip_inference = (mean_diff < g_config.motion_threshold);
}
prev_gray = gray;
if (!skip_inference && video_writer) video_writer->note_motion();
} else {
if (video_writer) video_writer->note_motion();
}
std::vector<Detection> detections;
float inf_ms = 0.0f;
if (!skip_inference) {
auto inf_start = std::chrono::steady_clock::now();
detections = model(frame);
inf_ms = std::chrono::duration<float, std::milli>(
std::chrono::steady_clock::now() - inf_start).count();
}
inf_ema = inf_ema * 0.9f + inf_ms * 0.1f;
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());
if (na > 0) store.ping_activity();
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) {
if (video_writer) video_writer->note_ayam_detection();
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, talenan_scores_vec(di))) {
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);
std::cout << "[" << now_str() << "] Batch #" << store.current_batch_number() << " | count: " << count << " | conf: " << ayam_scores_vec(di) << std::endl;
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;
crossing_times.push_back(mono);
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};
}
if (video_writer) video_writer->feed_frame(frame);
// 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();
while (!crossing_times.empty() &&
std::chrono::duration<double>(mono - crossing_times.front()).count() > g_config.rate_window_sec)
crossing_times.pop_front();
double rate = crossing_times.empty() ? 0.0 :
(crossing_times.size() / static_cast<double>(g_config.rate_window_sec) * 60.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_sec, rate);
draw_batch_banner(frame, w, batch_num, batch_pulse);
draw_footer(frame, w, h, frame_idx, is_live ? "LIVE" : "FILE", inf_ema, model_name);
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());
cv::Mat out_frame = frame;
if (blank_live_frame) {
out_frame = cv::Mat::zeros(frame.size(), frame.type());
blank_live_frame = false;
}
std::vector<uint8_t> jpeg;
cv::imencode(".jpg", out_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++;
prune_stale_tracks(ayam_tracked, g_config.tracked_prune_sec);
prune_stale_tracks(talenan_tracked, g_config.tracked_prune_sec);
}
cap.release();
if (video_writer) video_writer->shutdown();
store.shutdown();
std::cout << "\n=== Batch Summary (SQLite) ===" << std::endl;
std::cout << "Database: " << g_config.db_path << std::endl;
return 0;
}
View Git Blame Copy Permalink