sort.cpp

#include <iostream>
#include <fstream>
#include <iomanip> 
#include <unistd.h>    
#include <set>
#include "Hungarian.h"
#include "KalmanTracker.h"
#include <opencv2/highgui.hpp>
#include <opencv2/highgui/highgui_c.h>
#include <vector>
#include "opencv2/video/tracking.hpp"
#include "opencv2/highgui/highgui.hpp"
#include <gflags/gflags.h>
#include <filesystem>  // C++17 提供的文件系统库
#include <regex>
#include <chrono>
namespace fs = std::filesystem;
using namespace std;
using namespace cv;
DEFINE_string(input_path, "", "Path to the directory containing sequence subfolders");
DEFINE_string(img_ext, ".jpg", "图像文件扩展名，例如 .jpg, .bmp");
DEFINE_string(output_path, "output", "Directory to save result .txt and .avi files"); // 新增：输出路径参数
DEFINE_int32(max_age, 30, "Maximum number of frames to keep a track alive without new detections.");
DEFINE_int32(min_hits, 3, "Minimum number of associated detections before a track is displayed.");
DEFINE_double(iou_threshold, 0.3, "Minimum IOU for a match between a track and a detection.");// 新增：将SORT算法的关键参数暴露出来
typedef struct TrackingBox
{
	int frame;
	int id;
	Rect_<float> box;
}
TrackingBox;
    
// Computes IOU between two bounding boxes
double GetIOU(Rect_<float> bb_test, Rect_<float> bb_gt)
{
	float in = (bb_test & bb_gt).area();
	float un = bb_test.area() + bb_gt.area() - in;

	if (un < DBL_EPSILON)
		return 0;

	return (double)(in / un);
}
 
// global variables for counting
#define CNUM 20
int total_frames = 0;
double total_time = 0.0;

void TestSORT(string seqName, bool display, bool savevideo);

int main(int argc, char** argv)
{
	// 1. 初始化 GFLAGS
	gflags::ParseCommandLineFlags(&argc, &argv, true);

	// 2. 检查参数
	if (FLAGS_input_path.empty()) 
	{
		std::cerr << "Error: --input_path must be provided." << std::endl;
		return 1;
	}

	// 3. 遍历输入路径下所有子目录，作为序列名
	vector<string> sequences;
	for (const auto& entry : fs::directory_iterator(FLAGS_input_path)) 
	{
		if (entry.is_directory()) 
		{
			sequences.push_back(entry.path().filename().string());
		}
	}

	if (sequences.empty()) 
	{
		std::cerr << "No sequence folders found in path: " << FLAGS_input_path << std::endl;
		return 1;
	}
    auto total_start = std::chrono::high_resolution_clock::now();
	// 4. 执行每个序列的 tracking
	for (const auto& seq : sequences)
	{
		std::cout << "Processing sequence: " << seq << std::endl;
        TestSORT(seq,false, true);
    }
    
    // 5. 计算总耗时
    auto total_end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = total_end - total_start;
    total_time = elapsed.count();

    // 6. 安全输出结果
    if (total_time > 0 && total_frames > 0) 
	{
        double fps = static_cast<double>(total_frames) / total_time;
        std::cout << "\nTotal Tracking took: " << total_time << " seconds for "
                  << total_frames << " frames or " << fps << " FPS" << std::endl;
    } 
	else if (total_frames > 0) 
	{
        std::cout << "\nTotal Tracking took: ? seconds for "
                  << total_frames << " frames (time measurement failed)" << std::endl;
    } 
	else 
	{
        std::cout << "\nError: No frames processed. Check input data and paths." << std::endl;
    }

    return 0;
}
    

DEFINE_string(imgPath, "", "Path to the image folder");


void TestSORT(string seqName, bool display, bool savevideo)
{
        cout << "Processing " << seqName << "..." << endl;


	// 0. randomly generate colors, only for display
    RNG rng(0xFFFFFFFF);
    Scalar_<int> randColor[10];  // 假设 CNUM = 10
    for (int i = 0; i < 10; i++)
        rng.fill(randColor[i], RNG::UNIFORM, 0, 256);

    string imgPath = FLAGS_input_path + "/" + seqName + "/";

    if (display)
    {
        if (access(imgPath.c_str(), 0) == -1)
        {
            cerr << "Image path not found!" << endl;
            display = false;
        }
    }

    // 继续处理逻辑...
        VideoWriter videoWriter;
        int fps = 30.0;  
        int frame_width = 640;
        int frame_height = 380;
        bool videoInitialized = false;

	// 1. read detection file
	std::ifstream detectionFile;
    std::string detFileName =FLAGS_input_path + "/" + seqName + "/det.txt"; 
	detectionFile.open(detFileName);

	if (!detectionFile.is_open())
	{
		cerr << "Error: can not find file " << detFileName << endl;
		return;
	}

	string detLine;
	istringstream ss;
	vector<TrackingBox> detData;
	char ch;
	float tpx, tpy, tpw, tph;

	while ( getline(detectionFile, detLine) )
	{
		TrackingBox tb;

		ss.str(detLine);
		ss >> tb.frame >> ch >> tb.id >> ch;
		ss >> tpx >> ch >> tpy >> ch >> tpw >> ch >> tph;
		ss.str("");

		tb.box = Rect_<float>(Point_<float>(tpx, tpy), Point_<float>(tpx + tpw, tpy + tph));
		detData.push_back(tb);
	}
	detectionFile.close();

	// 2. group detData by frame
	int maxFrame = 0;
	for (auto tb : detData) // find max frame number
	{
		if (maxFrame < tb.frame)
			maxFrame = tb.frame;
	}

	vector<vector<TrackingBox>> detFrameData;
	vector<TrackingBox> tempVec;
	for (int fi = 0; fi < maxFrame; fi++)
	{
		for (auto tb : detData)
			if (tb.frame == fi + 1) // frame num starts from 1
				tempVec.push_back(tb);
		detFrameData.push_back(tempVec);
		tempVec.clear();
	}

	// 3. update across frames
	int frame_count = 0;
	//int max_age = 30;
	//int min_hits = 3;
	//double iouThreshold = 0.3;
	int max_age = FLAGS_max_age;
	int min_hits = FLAGS_min_hits;
    double iouThreshold = FLAGS_iou_threshold;
	vector<KalmanTracker> trackers;
	KalmanTracker::kf_count = 0; // tracking id relies on this, so we have to reset it in each seq.

	// variables used in the for-loop
	vector<Rect_<float>> predictedBoxes;
	vector<vector<double>> iouMatrix;
	vector<int> assignment;
	set<int> unmatchedDetections;
	set<int> unmatchedTrajectories;
	set<int> allItems;
	set<int> matchedItems;
	vector<cv::Point> matchedPairs;
	vector<TrackingBox> frameTrackingResult;
	unsigned int trkNum = 0;
	unsigned int detNum = 0;

	double cycle_time = 0.0;
	int64 start_time = 0;
	// prepare result file.
	fs::create_directories(FLAGS_output_path); 
	    ofstream resultsFile;
	//string resFileName = "output/" + seqName + ".txt";
	string resFileName = fs::path(FLAGS_output_path) / (seqName + ".txt");
	    resultsFile.open(resFileName);

	if (!resultsFile.is_open())
	{
		cerr << "Error: can not create file " << resFileName << endl;
		return;
	}

	//////////////////////////////////////////////
	// main loop
	for (int fi = 0; fi < maxFrame; fi++)
	{
		total_frames++;
		frame_count++;
		//cout << frame_count << endl;

		// I used to count running time using clock(), but found it seems to conflict with cv::cvWaitkey(),
		// when they both exists, clock() can not get right result. Now I use cv::getTickCount() instead.
		start_time = getTickCount();

		if (trackers.size() == 0) // the first frame met
		{
			// initialize kalman trackers using first detections.
			for (unsigned int i = 0; i < detFrameData[fi].size(); i++)
			{
				KalmanTracker trk = KalmanTracker(detFrameData[fi][i].box);
				trackers.push_back(trk);
			}
			// output the first frame detections
			for (unsigned int id = 0; id < detFrameData[fi].size(); id++)
			{
				TrackingBox tb = detFrameData[fi][id];
				resultsFile << tb.frame << "," << id + 1 << "," << tb.box.x << "," << tb.box.y << "," << tb.box.width << "," << tb.box.height << ",1,-1,-1,-1" << endl;
			}
			
		}

		///////////////////////////////////////
		// 3.1. get predicted locations from existing trackers.
		predictedBoxes.clear();

		for (auto it = trackers.begin(); it != trackers.end();)
		{
			Rect_<float> pBox = (*it).predict();
			if (pBox.x >= 0 && pBox.y >= 0)
			{
				predictedBoxes.push_back(pBox);
				it++;
			}
			else
			{
				it = trackers.erase(it);
				//cerr << "Box invalid at frame: " << frame_count << endl;
			}
		}

		///////////////////////////////////////
		// 3.2. associate detections to tracked object (both represented as bounding boxes)
		// dets : detFrameData[fi]
		trkNum = predictedBoxes.size();
		detNum = detFrameData[fi].size();

		iouMatrix.clear();
		iouMatrix.resize(trkNum, vector<double>(detNum, 0));

		for (unsigned int i = 0; i < trkNum; i++) // compute iou matrix as a distance matrix
		{
			for (unsigned int j = 0; j < detNum; j++)
			{
				// use 1-iou because the hungarian algorithm computes a minimum-cost assignment.
				iouMatrix[i][j] = 1 - GetIOU(predictedBoxes[i], detFrameData[fi][j].box);
			}
		}

		// solve the assignment problem using hungarian algorithm.
		// the resulting assignment is [track(prediction) : detection], with len=preNum
		HungarianAlgorithm HungAlgo;
		assignment.clear();
		HungAlgo.Solve(iouMatrix, assignment);

		// find matches, unmatched_detections and unmatched_predictions
		unmatchedTrajectories.clear();
		unmatchedDetections.clear();
		allItems.clear();
		matchedItems.clear();

		if (detNum > trkNum) //	there are unmatched detections
		{
			for (unsigned int n = 0; n < detNum; n++)
				allItems.insert(n);

			for (unsigned int i = 0; i < trkNum; ++i)
				matchedItems.insert(assignment[i]);

			set_difference(allItems.begin(), allItems.end(),
				matchedItems.begin(), matchedItems.end(),
				insert_iterator<set<int>>(unmatchedDetections, unmatchedDetections.begin()));
		}
		else if (detNum < trkNum) // there are unmatched trajectory/predictions
		{
			for (unsigned int i = 0; i < trkNum; ++i)
				if (assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm
					unmatchedTrajectories.insert(i);
		}
		else;

		// filter out matched with low IOU
		matchedPairs.clear();
		for (unsigned int i = 0; i < trkNum; ++i)
		{
			if (assignment[i] == -1) // pass over invalid values
				continue;
			if (1 - iouMatrix[i][assignment[i]] < iouThreshold)
			{
				unmatchedTrajectories.insert(i);
				unmatchedDetections.insert(assignment[i]);
			}
			else
				matchedPairs.push_back(cv::Point(i, assignment[i]));
		}

		///////////////////////////////////////
		// 3.3. updating trackers

		// update matched trackers with assigned detections.
		// each prediction is corresponding to a tracker
		int detIdx, trkIdx;
		for (unsigned int i = 0; i < matchedPairs.size(); i++)
		{
			trkIdx = matchedPairs[i].x;
			detIdx = matchedPairs[i].y;
			trackers[trkIdx].update(detFrameData[fi][detIdx].box);
		}

		// create and initialise new trackers for unmatched detections
		for (auto umd : unmatchedDetections)
		{
			KalmanTracker tracker = KalmanTracker(detFrameData[fi][umd].box);
			trackers.push_back(tracker);
		}

		// get trackers' output
		frameTrackingResult.clear();
		for (auto it = trackers.begin(); it != trackers.end();)
		{
			if (((*it).m_time_since_update < 1) &&
				((*it).m_hit_streak >= min_hits || frame_count <= min_hits))
			{
				TrackingBox res;
				res.box = (*it).get_state();
				res.id = (*it).m_id + 1;
				res.frame = frame_count;
				frameTrackingResult.push_back(res);
				it++;
			}
			else
				it++;

			// remove dead tracklet
			if (it != trackers.end() && (*it).m_time_since_update > max_age)
				it = trackers.erase(it);
		}

		cycle_time = (double)(getTickCount() - start_time);
		total_time += cycle_time / getTickFrequency();
		for (auto tb : frameTrackingResult)
			resultsFile << tb.frame << "," << tb.id << "," << tb.box.x << "," << tb.box.y << "," << tb.box.width << "," << tb.box.height << ",1,-1,-1,-1" << endl;

        bool needToProcessImage = display || savevideo;
        if (needToProcessImage)
        {
		
            ostringstream oss;
			oss << imgPath << setw(6) << setfill('0') << fi + 1 << FLAGS_img_ext;
			Mat img = imread(oss.str()); 
            
            if (img.empty()) 
			{
                // 如果找不到图像，可以打印警告并跳过，或创建一个黑屏帧
                cerr << "Warning: Cannot read image: " << oss.str() << endl;
                // 如果希望视频连续，可以创建一个黑屏
                if (videoInitialized) 
				{
                    Mat black_img = Mat::zeros(Size(frame_width, frame_height), CV_8UC3);
                    videoWriter.write(black_img);
                }
                continue;
            }
			for (auto tb : frameTrackingResult)
			{
			    cv::rectangle(img, tb.box, randColor[tb.id % CNUM], 2, 8, 0);
			    string label = to_string(tb.id);
                Point text_origin(tb.box.x, tb.box.y - 5); // 在框上方显示
    
                // 如果框在图像顶部，调整文本位置到框内
                if (text_origin.y < 10) 
                    text_origin.y = tb.box.y + 20;
    
                // 设置文本样式
                int fontFace = FONT_HERSHEY_SIMPLEX;
                double fontScale = 0.5;
                int thickness = 2;
    
                // 绘制文本背景（提高可读性）
                Size textSize = getTextSize(label, fontFace, fontScale, thickness, 0);
                rectangle(img, 
                    Point(text_origin.x - 1, text_origin.y - textSize.height - 2),
                    Point(text_origin.x + textSize.width + 1, text_origin.y + 2),
                    Scalar(0, 0, 0), FILLED);
    
                // 绘制ID文本
                putText(img, label, text_origin, fontFace, fontScale, 
                    Scalar(255, 255, 255), thickness);
				}
			if (savevideo)
			{
                if (!videoInitialized)
                {
                    frame_width = img.cols;
                    frame_height = img.rows;
                    //string outputVideoPath = "output/output_" + seqName + ".avi";
					string outputVideoPath = fs::path(FLAGS_output_path) / ("output_" + seqName + ".avi");
                    cout << "Initializing video writer to: " << outputVideoPath << endl;
                    videoWriter.open(outputVideoPath, VideoWriter::fourcc('M', 'J', 'P', 'G'), 30.0, Size(frame_width, frame_height));
                    
                    if (videoWriter.isOpened()) 
					{
                        videoInitialized = true;
                    } 
					else 
					{
                        cerr << "Error: Could not open video writer for " << outputVideoPath << endl;
                        savevideo = false; // 如果打开失败，停止后续的保存尝试
                    }
                }

                if (videoInitialized)
                {
                    videoWriter.write(img);
                }				
			}
			            if (display)
            {
                imshow(seqName, img);
                cvWaitKey(1); // 使用 1ms 延迟以允许GUI刷新
            }

			
		}
	}

	resultsFile.close();

	if (display)
	    destroyAllWindows();

}