91 lines
3.8 KiB
C
91 lines
3.8 KiB
C
|
#include "ANSLIB.h"
|
||
|
|
#define RETAINFRAMES 80
|
||
|
|
#define FILTERFRAMES 10
|
||
|
|
|
||
|
|
class CUSTOM_API ANSCustomWD : public IANSCustomClass
|
||
|
|
{
|
||
|
|
struct ImageSection {
|
||
|
|
cv::Rect region;
|
||
|
|
int priority;
|
||
|
|
ImageSection(const cv::Rect& r) : region(r), priority(0) {}
|
||
|
|
};
|
||
|
|
private:
|
||
|
|
using ANSLIBPtr = std::unique_ptr<ANSCENTER::ANSLIB, decltype(&ANSCENTER::ANSLIB::Destroy)>;
|
||
|
|
ANSLIBPtr _detector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
|
||
|
|
ANSLIBPtr _filter{ nullptr, &ANSCENTER::ANSLIB::Destroy };
|
||
|
|
int engineType;
|
||
|
|
std::string _detectorModelName;
|
||
|
|
std::string _detectorClassName;
|
||
|
|
int _detectorModelType; // Assuming 4 represents TensorRT YoloV11
|
||
|
|
int _detectorDetectionType; // Assuming 1 represents object detection
|
||
|
|
|
||
|
|
std::string _filterModelName;
|
||
|
|
std::string _filterClassName;
|
||
|
|
std::string _filterLabelMap;
|
||
|
|
|
||
|
|
int _filterModelType; // Assuming 1 represents OpenVINO YoloV12
|
||
|
|
int _filterDetectionType; // Assuming 1 represents object detection
|
||
|
|
|
||
|
|
std::recursive_mutex _mutex;
|
||
|
|
cv::Rect _detectedArea;// Area where weapon is detected
|
||
|
|
int _retainDetectedArea{ 0 };
|
||
|
|
bool _isWeaponDetected{ false };
|
||
|
|
float _detectionScoreThreshold{ 0.5 };
|
||
|
|
cv::Size previousImageSize = cv::Size(0, 0);
|
||
|
|
std::vector<ImageSection> cachedSections;
|
||
|
|
int _currentPriority{ 0 }; // None
|
||
|
|
int _realWeaponCheck{ 0 };
|
||
|
|
bool _isRealWeaponFrame{ false };
|
||
|
|
void UpdateNoDetectionCondition();
|
||
|
|
std::vector<ANSCENTER::Object> RunFilterGetPersons(const cv::Mat& frame);
|
||
|
|
// Function to seperate screen size
|
||
|
|
double calculateDistanceToCenter(const cv::Point& center, const cv::Rect& rect);
|
||
|
|
std::vector<ImageSection> divideImage(const cv::Mat& image);
|
||
|
|
int getHighestPriorityRegion();
|
||
|
|
int getLowestPriorityRegion();
|
||
|
|
cv::Rect getRegionByPriority(int priority);
|
||
|
|
|
||
|
|
// Utilities
|
||
|
|
float calculateIoU(const cv::Rect& box1, const cv::Rect& box2);
|
||
|
|
bool IsOverlapping(const ANSCENTER::Object& obj, const std::vector<ANSCENTER::Object>& objectList, float iouThreshold);
|
||
|
|
void UpdateActiveROI(const cv::Mat& frame, ANSCENTER::Object detectedObj);
|
||
|
|
std::vector<ANSCENTER::Object> ProcessExistingDetectedArea(
|
||
|
|
const cv::Mat& frame, const std::string& camera_id
|
||
|
|
#ifdef FNS_DEBUG
|
||
|
|
, cv::Mat& draw
|
||
|
|
#endif
|
||
|
|
);
|
||
|
|
bool ProcessWeaponDetection(
|
||
|
|
const cv::Mat& frame,
|
||
|
|
ANSCENTER::Object& detectedObj,
|
||
|
|
std::vector<ANSCENTER::Object>& output,
|
||
|
|
const std::vector<ANSCENTER::Object>& personDetections,
|
||
|
|
bool filterHadResults
|
||
|
|
#ifdef FNS_DEBUG
|
||
|
|
, cv::Mat& draw
|
||
|
|
#endif
|
||
|
|
);
|
||
|
|
void AddConfirmedWeaponDetection(
|
||
|
|
const cv::Mat& frame,
|
||
|
|
ANSCENTER::Object& detectedObj,
|
||
|
|
std::vector<ANSCENTER::Object>& output);
|
||
|
|
|
||
|
|
void ProcessNewDetectedArea(
|
||
|
|
const cv::Mat& frame,
|
||
|
|
const std::string& camera_id,
|
||
|
|
std::vector<ANSCENTER::Object>& output
|
||
|
|
#ifdef FNS_DEBUG
|
||
|
|
, cv::Mat& draw
|
||
|
|
#endif
|
||
|
|
);
|
||
|
|
void UpdateDetectedAreaFromObject(const cv::Mat& frame, const ANSCENTER::Object& detectedObj);
|
||
|
|
public:
|
||
|
|
bool Initialize(const std::string& modelDiretory,float detectionScoreThreshold, std::string& labelMap)override;
|
||
|
|
bool OptimizeModel(bool fp16)override;
|
||
|
|
std::vector<CustomObject> RunInference(const cv::Mat& input)override;
|
||
|
|
std::vector<CustomObject> RunInference(const cv::Mat& input, const std::string& camera_id)override;
|
||
|
|
bool Destroy()override;
|
||
|
|
virtual bool ConfigureParameters(CustomParams& param) override;
|
||
|
|
ANSCustomWD();
|
||
|
|
~ANSCustomWD();
|
||
|
|
};
|