AbandonedObjectionsAlgo/demo/StraDemo/StraDemo.cpp

#include <chrono>
#include <fstream>
#include <opencv2/opencv.hpp>

#include "Logger.h"
#include "OSTime.h"
#include "SysUtils.h"
#include "Semaphore.h"
#include "MemPool.h"
#include "DetUtils.h"

#include "DetectorAPI.h"
#include "AbandObjDetector.h"
#include "ClipDetector.h"

#include "StrategyAPI.h"
#include "AbandClipStrategy.h"

struct DetectionResult
{
    cv::Rect rect;
    std::string label;
    std::chrono::steady_clock::time_point timestamp;
};

struct CallbackContext
{
    masd::Strategy *strategy;
    cv::Mat *img;
    cv::VideoWriter *videoWriter;
    std::vector<DetectionResult> detections;
};

CallbackContext callbackContext;
tzc::Semaphore SEMA;

void SaveDetectorInfo(const char *info)
{
    std::ofstream outFile("detector_info.json", std::ios::out | std::ios::app);
    if (outFile.is_open())
    {
        outFile << "{\n\t\"Detector Information\": \"" << info << "\"\n}" << std::endl;
        outFile.close();
        TZLogInfo("Detector information saved to detector_info.json~~~");
    }
    else
    {
        TZLogError("Failed to open file for writing!!!");
    }
}

TZ_INT DetectionCallback(SPtr<masd::StreamInfo> &media, void *ctx)
{
    CallbackContext *callbackContext = reinterpret_cast<CallbackContext *>(ctx);
    if (!callbackContext || !callbackContext->img || callbackContext->strategy == nullptr)
        return -1;

    cv::Mat *img = callbackContext->img;
    masd::Strategy *strategy = callbackContext->strategy;
    if (img->empty())
        return -1;

    auto allDetRst = media->GetAllDetRst();
    std::vector<DetectionResult> currentDetections;

    for (const auto &detPair : allDetRst)
    {
        const std::string &detKey = detPair.first;
        const auto &detProducing = detPair.second;
        cv::Rect bbox;
        std::string label;
        if (detKey == "Clip")
        {
            masd::clip::ClipDetectResult clipResult;
            if (masd::clip::ClipDetectResult::fromJson(detProducing->Result, clipResult))
            {
                for (const auto &box : clipResult.boxes)
                    bbox = cv::Rect(box.x, box.y, box.w, box.h);
            }
            strategy->DoStrategy(media);
            SPtr<masd::StraProducing> straProducing = media->GetStraProducing();
            label = straProducing ? straProducing->Result.RstName : "Unknown";
        }
        currentDetections.push_back({bbox, label, std::chrono::steady_clock::now()});
    }

    auto &cachedDetections = callbackContext->detections;
    const auto now = std::chrono::steady_clock::now();
    const std::chrono::seconds timeout(4);
    std::vector<DetectionResult> updatedDetections;
    for (const auto &detection : cachedDetections)
    {
        if (now - detection.timestamp < timeout)
            updatedDetections.push_back(detection);
    }
    updatedDetections.insert(updatedDetections.end(), currentDetections.begin(), currentDetections.end());
    callbackContext->detections = std::move(updatedDetections);

    for (const auto &detection : callbackContext->detections)
    {
        cv::rectangle(*img, detection.rect, cv::Scalar(0, 255, 0), 2);
        cv::putText(*img, detection.label, detection.rect.tl() - cv::Point(0, 10),
                    cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);
    }

    SEMA.Signal();
    return 0;
}

int main()
{
    INITIALIZE_LOGGER_NORMAL("test", "./test.log", 1, 100, 6, 1, 1);
    masd::MemPool *pool = masd::MEMPOOL;
    if (pool->Initialize() != masd::MEC_OK)
        return -1;

    masd::Detector *abandObjDet = new masd::abandobj::AbandObjDetector("AbandObj", "遗留物检测");
    if (abandObjDet == nullptr)
    {
        Dispose();
        return -1;
    }

    masd::Detector *clipDet = new masd::clip::ClipDetector("Clip", "Clip检测");
    if (clipDet == nullptr)
    {
        Dispose();
        return -1;
    }

    masd::Strategy *strategy = new masd::abandclip::AbandClipStrategy();
    if (strategy == nullptr)
    {
        Dispose();
        return -1;
    }

    std::string initAbandObjParam = "{\"algo\": \"MOG2\", \"factor\": 4, \"short_term_rate\": 0.01, \"short_term_history\": 200, \"long_term_rate\": 0.0005, \"long_term_history\": 5000, \"var_threshold\": 16.0, \"detect_shadows\": true, \"iou_threshold\": 0.6, \"area_threshold\": 625.0, \"perimeter_threshold\": 100.0}";
    if (abandObjDet->Initialize(initAbandObjParam) != masd::MEC_OK)
        return -1;

    std::string initClipParam = "{\"gpu_id\": 0, \"model_path\": \"/home/flechazo/workspace/abandoned-object-detection-algo/clip/models/visual.onnx\"}";
    if (clipDet->Initialize(initClipParam) != masd::MEC_OK)
        return -1;

    if (strategy->Initialize() != masd::MEC_OK)
        return -1;

    std::string AbandClipConfig = "{\"FeatureThreshold\": -100.0, \"FeatureFilePath\": \"../textfeatures.json\"}";
    if (strategy->SetStrategyCfg(AbandClipConfig) != masd::MEC_OK)
        return -1;

    cv::VideoCapture videoCapture("../../media/video.mp4");
    if (!videoCapture.isOpened())
        return -1;

    cv::VideoWriter videoWriter("output_video.mp4",
                                cv::VideoWriter::fourcc('X', '2', '6', '4'),
                                static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FPS)),
                                cv::Size(static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FRAME_WIDTH)),
                                         static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FRAME_HEIGHT))));
    if (!videoWriter.isOpened())
        return -1;

    cv::Mat frame;
    while (videoCapture.read(frame))
    {
        if (frame.empty())
            break;

        TZ_INT length = frame.total() * frame.elemSize();
        void *allocatedMem = pool->AllocAvailMem(length);
        if (!allocatedMem)
            break;

        std::memset(allocatedMem, 0, length);
        SPtr<masd::Media> mediaResource = std::make_shared<masd::Media>(length);
        mediaResource->Width = frame.cols;
        mediaResource->Height = frame.rows;
        mediaResource->DataType = frame.type();
        mediaResource->Mem = allocatedMem;
        std::memcpy(mediaResource->Mem, frame.data, length);

        SPtr<masd::StreamInfo> streamInfo = std::make_shared<masd::StreamInfo>();
        streamInfo->SetMediaRsc(mediaResource);

        callbackContext.img = new cv::Mat(frame.clone());
        callbackContext.videoWriter = &videoWriter;
        callbackContext.strategy = strategy;

        abandObjDet->DoDetect(streamInfo, DetectionCallback, &callbackContext);
        SEMA.Wait();
        clipDet->DoDetect(streamInfo, DetectionCallback, &callbackContext);
        SEMA.Wait();

        videoWriter.write(*callbackContext.img);

        pool->ReturnFreeMem(allocatedMem, length);
        delete callbackContext.img;
        callbackContext.img = nullptr;
    }

    char detectorInfo[65536];
    if (GetInformation(detectorInfo) != masd::MEC_OK)
    {
        Dispose();
        return -1;
    }
    SaveDetectorInfo(detectorInfo);

    TZ_delete(abandObjDet);
    TZ_delete(clipDet);
    TZ_delete(strategy);

    if (Dispose() != masd::MEC_OK)
        return -1;

    return 0;
}