#include "ANSTENSORRTPOSE.h" #include "Utility.h" #include #include namespace ANSCENTER { bool ANSTENSORRTPOSE::OptimizeModel(bool fp16, std::string& optimizedModelFolder) { std::lock_guard lock(_mutex); if (!ANSODBase::OptimizeModel(fp16, optimizedModelFolder)) { return false; } if (!FileExist(_modelFilePath)) { this->_logger.LogFatal("ANSTENSORRTPOSE::OptimizeModel", "Raw model file path does not exist", __FILE__, __LINE__); return false; } try { _fp16 = fp16; optimizedModelFolder = GetParentFolder(_modelFilePath); // Check if the engine already exists to avoid reinitializing if (!m_trtEngine) { // Fixed batch size of 1 for this model m_options.optBatchSize = _modelConfig.gpuOptBatchSize; m_options.maxBatchSize = _modelConfig.gpuMaxBatchSize; m_options.deviceIndex = _modelConfig.gpuDeviceIndex; m_options.maxInputHeight = _modelConfig.maxInputHeight; m_options.minInputHeight = _modelConfig.minInputHeight; m_options.optInputHeight = _modelConfig.optInputHeight; m_options.maxInputWidth = _modelConfig.maxInputWidth; m_options.minInputWidth = _modelConfig.minInputWidth; m_options.optInputWidth = _modelConfig.optInputWidth; m_options.engineFileDir = optimizedModelFolder; m_options.precision = (_fp16 ? Precision::FP16 : Precision::FP32); // Create the TensorRT inference engine m_trtEngine = std::make_unique>(m_options); } // Build the TensorRT engine auto succ = m_trtEngine->buildWithRetry(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE); if (!succ) { const std::string errMsg = "Error: Unable to build the TensorRT engine. " "Try increasing TensorRT log severity to kVERBOSE."; this->_logger.LogError("ANSTENSORRTPOSE::OptimizeModel", errMsg, __FILE__, __LINE__); _modelLoadValid = false; return false; } _modelLoadValid = true; return true; } catch (const std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::OptimizeModel", e.what(), __FILE__, __LINE__); optimizedModelFolder.clear(); return false; } } bool ANSTENSORRTPOSE::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) { std::lock_guard lock(_mutex); try { bool result = ANSODBase::LoadModel(modelZipFilePath, modelZipPassword); if (!result) return false; _modelConfig.detectionType = ANSCENTER::DetectionType::DETECTION; _modelConfig.modelType = ModelType::TENSORRT; _modelConfig.inpHeight = 640; _modelConfig.inpWidth = 640; if (_modelConfig.modelMNSThreshold < 0.2) _modelConfig.modelMNSThreshold = 0.5; if (_modelConfig.modelConfThreshold < 0.2) _modelConfig.modelConfThreshold = 0.5; if (_modelConfig.numKPS <= 0 || _modelConfig.numKPS > 133) // 133 = COCO wholebody max _modelConfig.numKPS = 17; if (_modelConfig.kpsThreshold == 0)_modelConfig.kpsThreshold = 0.5; // If not define _fp16 = true; // Load Model from Here // Load Model from Here TOP_K = 100; SEG_CHANNELS = 32; PROBABILITY_THRESHOLD = _modelConfig.detectionScoreThreshold; NMS_THRESHOLD = _modelConfig.modelMNSThreshold; SEGMENTATION_THRESHOLD = 0.5f; SEG_H = 160; SEG_W = 160; NUM_KPS = _modelConfig.numKPS; KPS_THRESHOLD = _modelConfig.kpsThreshold; SEG_CHANNELS = 32; // For segmentation if (!m_trtEngine) { // Fixed batch size of 1 for this model m_options.optBatchSize = _modelConfig.gpuOptBatchSize; m_options.maxBatchSize = _modelConfig.gpuMaxBatchSize; m_options.deviceIndex = _modelConfig.gpuDeviceIndex; m_options.maxInputHeight = _modelConfig.maxInputHeight; m_options.minInputHeight = _modelConfig.minInputHeight; m_options.optInputHeight = _modelConfig.optInputHeight; m_options.maxInputWidth = _modelConfig.maxInputWidth; m_options.minInputWidth = _modelConfig.minInputWidth; m_options.optInputWidth = _modelConfig.optInputWidth; m_options.engineFileDir = _modelFolder; // Use FP16 or FP32 precision based on the input flag m_options.precision = (_fp16 ? Precision::FP16 : Precision::FP32); // Create the TensorRT inference engine m_trtEngine = std::make_unique>(m_options); } // 0. Check if the configuration file exist if (FileExist(_modelConfigFile)) { ModelType modelType; std::vector inputShape; _classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape); if (inputShape.size() == 2) { if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0]; if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1]; } } else {// This is old version of model zip file _modelFilePath = CreateFilePath(_modelFolder, "train_last.onnx"); _classFilePath = CreateFilePath(_modelFolder, "classes.names"); std::ifstream isValidFileName(_classFilePath); if (!isValidFileName) { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__); LoadClassesFromString(); } else { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__); LoadClassesFromFile(); } } // Load the TensorRT engine file if (this->_loadEngineOnCreation) { auto succ = m_trtEngine->buildLoadNetwork(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE, m_maxSlotsPerGpu); if (!succ) { const std::string errMsg = "Error: Unable to load TensorRT engine weights into memory. " + _modelFilePath; this->_logger.LogError("ANSTENSORRTPOSE::Initialize", errMsg, __FILE__, __LINE__); _modelLoadValid = false; return false; } } _modelLoadValid = true; _isInitialized = true; return true; } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::LoadModel", e.what(), __FILE__, __LINE__); return false; } } bool ANSTENSORRTPOSE::LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap) { std::lock_guard lock(_mutex); try { bool result = ANSODBase::LoadModelFromFolder(licenseKey, modelConfig, modelName, className, modelFolder, labelMap); if (!result) return false; _modelConfig = modelConfig; _modelConfig.detectionType = ANSCENTER::DetectionType::DETECTION; _modelConfig.modelType = ModelType::RTPOSE; _modelConfig.inpHeight = 640; _modelConfig.precisionType = PrecisionType::FP32; // Default to FP16 for TensorRT models if (_modelConfig.numKPS <= 0 || _modelConfig.numKPS > 133) // 133 = COCO wholebody max _modelConfig.numKPS = 17; _modelConfig.inpWidth = 640; if (_modelConfig.modelMNSThreshold < 0.2) _modelConfig.modelMNSThreshold = 0.5; if (_modelConfig.modelConfThreshold < 0.2) _modelConfig.modelConfThreshold = 0.5; if (_modelConfig.kpsThreshold <= 0)_modelConfig.kpsThreshold = 0.5; // If not define _fp16 = true; // Load Model from Here // Load Model from Here TOP_K = 100; SEG_CHANNELS = 32; PROBABILITY_THRESHOLD = _modelConfig.detectionScoreThreshold; NMS_THRESHOLD = _modelConfig.modelMNSThreshold; SEGMENTATION_THRESHOLD = 0.5f; SEG_H = 160; SEG_W = 160; NUM_KPS = _modelConfig.numKPS; KPS_THRESHOLD = _modelConfig.kpsThreshold; SEG_CHANNELS = 32; // For segmentation std::string _modelName = modelName; if (_modelName.empty()) { _modelName = "train_last"; } std::string modelFullName = _modelName + ".onnx"; if (!m_trtEngine) { // Fixed batch size of 1 for this model m_options.optBatchSize = _modelConfig.gpuOptBatchSize; m_options.maxBatchSize = _modelConfig.gpuMaxBatchSize; m_options.deviceIndex = _modelConfig.gpuDeviceIndex; m_options.maxInputHeight = _modelConfig.maxInputHeight; m_options.minInputHeight = _modelConfig.minInputHeight; m_options.optInputHeight = _modelConfig.optInputHeight; m_options.maxInputWidth = _modelConfig.maxInputWidth; m_options.minInputWidth = _modelConfig.minInputWidth; m_options.optInputWidth = _modelConfig.optInputWidth; m_options.engineFileDir = _modelFolder; // Use FP16 or FP32 precision based on the input flag m_options.precision = (_fp16 ? Precision::FP16 : Precision::FP32); // Create the TensorRT inference engine m_trtEngine = std::make_unique>(m_options); } // 0. Check if the configuration file exist if (FileExist(_modelConfigFile)) { ModelType modelType; std::vector inputShape; _classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape); if (inputShape.size() == 2) { if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0]; if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1]; } } else {// This is old version of model zip file _modelFilePath = CreateFilePath(_modelFolder, modelFullName); _classFilePath = CreateFilePath(_modelFolder, className); std::ifstream isValidFileName(_classFilePath); if (!isValidFileName) { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__); LoadClassesFromString(); } else { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__); LoadClassesFromFile(); } } // 1. Load labelMap and engine labelMap.clear(); if (!_classes.empty()) labelMap = VectorToCommaSeparatedString(_classes); // Load the TensorRT engine file if (this->_loadEngineOnCreation) { auto succ = m_trtEngine->buildLoadNetwork(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE, m_maxSlotsPerGpu); if (!succ) { const std::string errMsg = "Error: Unable to load TensorRT engine weights into memory. " + _modelFilePath; this->_logger.LogError("ANSTENSORRTPOSE::Initialize", errMsg, __FILE__, __LINE__); _modelLoadValid = false; return false; } } _modelLoadValid = true; _isInitialized = true; return true; } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::LoadModelFromFolder", e.what(), __FILE__, __LINE__); return false; } } bool ANSTENSORRTPOSE::Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) { std::lock_guard lock(_mutex); try { const bool engineAlreadyLoaded = _modelLoadValid && _isInitialized && m_trtEngine != nullptr; _modelLoadValid = false; bool result = ANSODBase::Initialize(licenseKey, modelConfig, modelZipFilePath, modelZipPassword, labelMap); if (!result) return false; // Parsing for YOLO only here _modelConfig = modelConfig; _modelConfig.detectionType = ANSCENTER::DetectionType::DETECTION; _modelConfig.modelType = ModelType::TENSORRT; _modelConfig.inpHeight = 640; _modelConfig.inpWidth = 640; if (_modelConfig.numKPS <= 0 || _modelConfig.numKPS > 133) // 133 = COCO wholebody max _modelConfig.numKPS = 17; _modelConfig.precisionType = PrecisionType::FP32; // Default to FP16 for TensorRT models if (_modelConfig.modelMNSThreshold < 0.2) _modelConfig.modelMNSThreshold = 0.5; if (_modelConfig.modelConfThreshold < 0.2) _modelConfig.modelConfThreshold = 0.5; if (_modelConfig.kpsThreshold == 0)_modelConfig.kpsThreshold = 0.5; // If not define _fp16 = true; // Load Model from Here // Load Model from Here TOP_K = 100; SEG_CHANNELS = 32; PROBABILITY_THRESHOLD = _modelConfig.detectionScoreThreshold; NMS_THRESHOLD = _modelConfig.modelMNSThreshold; SEGMENTATION_THRESHOLD = 0.5f; SEG_H = 160; SEG_W = 160; NUM_KPS = _modelConfig.numKPS; KPS_THRESHOLD = _modelConfig.kpsThreshold; SEG_CHANNELS = 32; // For segmentation if (!m_trtEngine) { // Fixed batch size of 1 for this model m_options.optBatchSize = _modelConfig.gpuOptBatchSize; m_options.maxBatchSize = _modelConfig.gpuMaxBatchSize; m_options.deviceIndex = _modelConfig.gpuDeviceIndex; m_options.maxInputHeight = _modelConfig.maxInputHeight; m_options.minInputHeight = _modelConfig.minInputHeight; m_options.optInputHeight = _modelConfig.optInputHeight; m_options.maxInputWidth = _modelConfig.maxInputWidth; m_options.minInputWidth = _modelConfig.minInputWidth; m_options.optInputWidth = _modelConfig.optInputWidth; m_options.engineFileDir = _modelFolder; // Use FP16 or FP32 precision based on the input flag m_options.precision = (_fp16 ? Precision::FP16 : Precision::FP32); // Create the TensorRT inference engine m_trtEngine = std::make_unique>(m_options); } // 0. Check if the configuration file exist if (FileExist(_modelConfigFile)) { ModelType modelType; std::vector inputShape; _classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape); if (inputShape.size() == 2) { if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0]; if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1]; } } else {// This is old version of model zip file _modelFilePath = CreateFilePath(_modelFolder, "train_last.onnx"); _classFilePath = CreateFilePath(_modelFolder, "classes.names"); std::ifstream isValidFileName(_classFilePath); if (!isValidFileName) { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__); LoadClassesFromString(); } else { this->_logger.LogDebug("ANSTENSORRTPOSE::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__); LoadClassesFromFile(); } } // 1. Load labelMap and engine labelMap.clear(); if (!_classes.empty()) labelMap = VectorToCommaSeparatedString(_classes); // Load the TensorRT engine file if (this->_loadEngineOnCreation && !engineAlreadyLoaded) { auto succ = m_trtEngine->buildLoadNetwork(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE, m_maxSlotsPerGpu); if (!succ) { const std::string errMsg = "Error: Unable to load TensorRT engine weights into memory. " + _modelFilePath; this->_logger.LogError("ANSTENSORRTPOSE::Initialize", errMsg, __FILE__, __LINE__); _modelLoadValid = false; return false; } } _modelLoadValid = true; _isInitialized = true; return true; } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::Initialize", e.what(), __FILE__, __LINE__); return false; } } std::vector ANSTENSORRTPOSE::RunInference(const cv::Mat& inputImgBGR) { return RunInference(inputImgBGR, ""); } std::vector ANSTENSORRTPOSE::RunInference(const cv::Mat& inputImgBGR,const std::string& camera_id) { { std::lock_guard lock(_mutex); // Validation checks if (!_modelLoadValid) { _logger.LogError("ANSTENSORRTPOSE::RunInference", "Cannot load the TensorRT model. Please check if it exists", __FILE__, __LINE__); return {}; } if (!_licenseValid) { _logger.LogError("ANSTENSORRTPOSE::RunInference", "Runtime license is not valid or expired. Please contact ANSCENTER", __FILE__, __LINE__); return {}; } if (!_isInitialized) { _logger.LogError("ANSTENSORRTPOSE::RunInference", "Model is not initialized", __FILE__, __LINE__); return {}; } if (inputImgBGR.empty() || inputImgBGR.cols < 10 || inputImgBGR.rows < 10) { return {}; } } try { return DetectObjects(inputImgBGR, camera_id); } catch (const std::exception& e) { _logger.LogFatal("ANSTENSORRTPOSE::RunInference", e.what(), __FILE__, __LINE__); return {}; } } ANSTENSORRTPOSE::~ANSTENSORRTPOSE() { try { Destroy(); } catch (std::exception& e) { this->_logger.LogError("ANSTENSORRTPOSE::~ANSTENSORRTPOSE()", e.what(), __FILE__, __LINE__); } } bool ANSTENSORRTPOSE::Destroy() { try { m_trtEngine.reset(); // Releases the current engine and sets m_trtEngine to nullptr. return true; } catch (std::exception& e) { this->_logger.LogError("ANSTENSORRTPOSE::~ANSTENSORRTPOSE()", e.what(), __FILE__, __LINE__); return false; } } // private std::vector ANSTENSORRTPOSE::DetectObjects(const cv::Mat& inputImage, const std::string& camera_id) { try { // --- 1. Set GPU device context --- if (m_trtEngine) { m_trtEngine->setDeviceContext(); } // --- 1b. CUDA context health check --- if (!m_nv12Helper.isCudaContextHealthy(_logger, "ANSTENSORRTPOSE")) { return {}; } // --- 2. Preprocess under lock --- // Try NV12 fast path first, falls back to standard GPU preprocessing. ImageMetadata meta; std::vector> input; bool usedNV12 = false; float bgrFullResScaleX = 1.0f, bgrFullResScaleY = 1.0f; { std::lock_guard lock(_mutex); const int inferenceGpu = m_trtEngine ? m_trtEngine->getPreferredDeviceIndex() : 0; const auto& inputDims = m_trtEngine->getInputDims(); const int inputW = inputDims[0].d[2]; const int inputH = inputDims[0].d[1]; auto nv12 = m_nv12Helper.tryNV12(inputImage, inferenceGpu, inputW, inputH, NV12PreprocessHelper::defaultYOLOLauncher(), _logger, "ANSTENSORRTPOSE"); if (nv12.succeeded) { meta.imgWidth = nv12.metaWidth; meta.imgHeight = nv12.metaHeight; meta.ratio = nv12.ratio; input = {{ std::move(nv12.gpuRGB) }}; usedNV12 = true; } else if (nv12.useBgrFullRes) { input = Preprocess(nv12.bgrFullResImg, meta); usedNV12 = !input.empty(); bgrFullResScaleX = nv12.bgrFullResScaleX; bgrFullResScaleY = nv12.bgrFullResScaleY; } if (input.empty()) { input = Preprocess(inputImage, meta); } m_nv12Helper.tickInference(); } if (input.empty()) return {}; // Phase 2: Inference — mutex released; pool dispatches to idle GPU slot std::vector>> featureVectors; auto succ = m_trtEngine->runInference(input, featureVectors); if (!succ) { this->_logger.LogError("ANSTENSORRTPOSE::DetectObjects", "Error running inference", __FILE__, __LINE__); return {}; } // Phase 3: Postprocess under lock std::lock_guard lock(_mutex); std::vector featureVector; Engine::transformOutput(featureVectors, featureVector); auto ret = PostProcessPose(featureVector, camera_id, meta); // --- 4b. Rescale coords from full-res to display-res (BGR full-res path) --- if (bgrFullResScaleX != 1.0f || bgrFullResScaleY != 1.0f) { for (auto& obj : ret) { obj.box.x = static_cast(obj.box.x * bgrFullResScaleX); obj.box.y = static_cast(obj.box.y * bgrFullResScaleY); obj.box.width = static_cast(obj.box.width * bgrFullResScaleX); obj.box.height = static_cast(obj.box.height * bgrFullResScaleY); for (auto& pt : obj.polygon) { pt.x *= bgrFullResScaleX; pt.y *= bgrFullResScaleY; } for (size_t k = 0; k + 2 < obj.kps.size(); k += 3) { obj.kps[k] *= bgrFullResScaleX; obj.kps[k + 1] *= bgrFullResScaleY; } } } if (_trackerEnabled) { ret = ApplyTracking(ret, camera_id); if (_stabilizationEnabled) ret = StabilizeDetections(ret, camera_id); } return ret; } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::DetectObjects", e.what(), __FILE__, __LINE__); return {}; } } std::vector> ANSTENSORRTPOSE::Preprocess(const cv::Mat& inputImage, ImageMetadata& outMeta) { try { if (!_licenseValid) { this->_logger.LogFatal("ANSTENSORRTPOSE::Preprocess", "Invalid license", __FILE__, __LINE__); return {}; } // Get model input dimensions const auto& inputDims = m_trtEngine->getInputDims(); const int inputH = inputDims[0].d[1]; const int inputW = inputDims[0].d[2]; // Upload the image to GPU memory cv::cuda::Stream stream; cv::cuda::GpuMat img; if (inputImage.channels() == 1) { // Convert grayscale to 3-channel BGR before uploading cv::Mat img3Channel; cv::cvtColor(inputImage, img3Channel, cv::COLOR_GRAY2BGR); img.upload(img3Channel, stream); } else { img.upload(inputImage, stream); } // Convert to RGB cv::cuda::GpuMat imgRGB; cv::cuda::cvtColor(img, imgRGB, cv::COLOR_BGR2RGB, 0, stream); stream.waitForCompletion(); // Set image size parameters outMeta.imgHeight = imgRGB.rows; outMeta.imgWidth = imgRGB.cols; if (outMeta.imgHeight > 0 && outMeta.imgWidth > 0) { outMeta.ratio = 1.f / std::min(inputDims[0].d[2] / static_cast(imgRGB.cols), inputDims[0].d[1] / static_cast(imgRGB.rows)); cv::cuda::GpuMat resized = imgRGB; // Resize to the model's expected input size while maintaining aspect ratio with padding if (resized.rows != inputDims[0].d[1] || resized.cols != inputDims[0].d[2]) { resized = Engine::resizeKeepAspectRatioPadRightBottom(imgRGB, inputDims[0].d[1], inputDims[0].d[2]); } // Convert to format expected by our inference engine std::vector input{ std::move(resized) }; std::vector> inputs{ std::move(input) }; return inputs; } else { this->_logger.LogFatal("TENSORRTCL::Preprocess", "Image height or width is zero after processing (Width: " + std::to_string(outMeta.imgWidth) + ", Height: " + std::to_string(outMeta.imgHeight) + ")", __FILE__, __LINE__); return {}; } } catch (const std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::Preprocess", e.what(), __FILE__, __LINE__); return {}; } } std::vector ANSTENSORRTPOSE::PostProcessPose(std::vector& featureVector, const std::string& camera_id, const ImageMetadata& meta) { try { const auto& outputDims = m_trtEngine->getOutputDims(); auto numChannels = outputDims[0].d[1]; auto numAnchors = outputDims[0].d[2]; std::vector bboxes; std::vector scores; std::vector labels; std::vector indices; std::vector>kpss; cv::Mat output = cv::Mat(numChannels, numAnchors, CV_32F, featureVector.data()); output = output.t(); // Get all the YOLO proposals for (int i = 0; i < numAnchors; i++) { auto rowPtr = output.row(i).ptr(); auto bboxesPtr = rowPtr; auto scoresPtr = rowPtr + 4; auto kps_ptr = rowPtr + 5; float score = *scoresPtr; if (score > this->_modelConfig.detectionScoreThreshold) { float x = *bboxesPtr++; float y = *bboxesPtr++; float w = *bboxesPtr++; float h = *bboxesPtr; float x0 = std::clamp((x - 0.5f * w) * meta.ratio, 0.f, meta.imgWidth); float y0 = std::clamp((y - 0.5f * h) * meta.ratio, 0.f, meta.imgHeight); float x1 = std::clamp((x + 0.5f * w) * meta.ratio, 0.f, meta.imgWidth); float y1 = std::clamp((y + 0.5f * h) * meta.ratio, 0.f, meta.imgHeight); cv::Rect_ bbox; bbox.x = x0; bbox.y = y0; bbox.width = x1 - x0; bbox.height = y1 - y0; bbox.x = std::max(0.f, bbox.x); bbox.y = std::max(0.f, bbox.y); bbox.width = std::min(meta.imgWidth - bbox.x, bbox.width); bbox.height = std::min(meta.imgHeight - bbox.y, bbox.height); std::vector kps; for (int k = 0; k < NUM_KPS; k++) { float kpsX = *(kps_ptr + 3 * k) * meta.ratio; float kpsY = *(kps_ptr + 3 * k + 1) * meta.ratio; float kpsS = *(kps_ptr + 3 * k + 2); kpsX = std::clamp(kpsX, 0.f, meta.imgWidth); kpsY = std::clamp(kpsY, 0.f, meta.imgHeight); cv::Point2f kp(kpsX, kpsY); kps.push_back(kp); } bboxes.push_back(bbox); labels.push_back(0); // All detected objects are people scores.push_back(score); kpss.push_back(kps); } } // Run NMS cv::dnn::NMSBoxesBatched(bboxes, scores, labels, PROBABILITY_THRESHOLD, NMS_THRESHOLD, indices); std::vector objects; int classNameSize = static_cast(_classes.size()); // Choose the top k detections for (auto& chosenIdx : indices) { if (scores[chosenIdx] > _modelConfig.detectionScoreThreshold) { std::stringstream keypointXss; std::stringstream keypointYss; std::vector keypoints; for (size_t keypointIdx = 0; keypointIdx < kpss[chosenIdx].size(); keypointIdx++) { keypointXss << kpss[chosenIdx][keypointIdx].x; keypointYss << kpss[chosenIdx][keypointIdx].y; // Add semicolon after each value except the last one if (keypointIdx < kpss[chosenIdx].size() - 1) { keypointXss << ";"; keypointYss << ";"; } keypoints.push_back(kpss[chosenIdx][keypointIdx].x); keypoints.push_back(kpss[chosenIdx][keypointIdx].y); } std::string keypointXString = keypointXss.str(); std::string keypointYString = keypointYss.str(); std::string keypointString = keypointXString + "|" + keypointYString; Object obj{}; obj.confidence = scores[chosenIdx]; obj.classId = labels[chosenIdx]; if (!_classes.empty()) { if (obj.classId < classNameSize) { obj.className = _classes[obj.classId]; } else { obj.className = _classes[classNameSize - 1]; // Use last valid class name if out of range } } else { obj.className = "Unknown"; // Fallback if _classes is empty } obj.box = bboxes[chosenIdx]; obj.polygon = kpss[chosenIdx]; obj.kps = keypoints; obj.cameraId = camera_id; obj.extraInfo = keypointString; objects.push_back(obj); } } //EnqueueDetection(objects, camera_id); return objects; } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::PostProcessPose", e.what(), __FILE__, __LINE__); std::vector result; result.clear(); return result; } } std::vector> ANSTENSORRTPOSE::DetectObjectsBatch(const std::vector& inputImages, const std::string& camera_id) { { std::lock_guard lock(_mutex); if (inputImages.empty()) { _logger.LogFatal("ANSTENSORRTPOSE::DetectObjectsBatch", "Empty input images vector", __FILE__, __LINE__); return {}; } } // Auto-split if batch exceeds engine capacity const int maxBatch = m_options.maxBatchSize > 0 ? m_options.maxBatchSize : 1; if (static_cast(inputImages.size()) > maxBatch) { const size_t numImages = inputImages.size(); std::vector> allResults; allResults.reserve(numImages); // Process chunks sequentially to avoid GPU contention on the same engine for (size_t start = 0; start < numImages; start += static_cast(maxBatch)) { const size_t end = std::min(start + static_cast(maxBatch), numImages); std::vector chunk(inputImages.begin() + start, inputImages.begin() + end); auto chunkResults = DetectObjectsBatch(chunk, camera_id); if (chunkResults.size() == chunk.size()) { for (auto& r : chunkResults) allResults.push_back(std::move(r)); } else { _logger.LogError("ANSTENSORRTPOSE::DetectObjectsBatch", "Chunk returned " + std::to_string(chunkResults.size()) + " results, expected " + std::to_string(chunk.size()) + ". Padding with empty results.", __FILE__, __LINE__); for (auto& r : chunkResults) allResults.push_back(std::move(r)); for (size_t pad = chunkResults.size(); pad < chunk.size(); ++pad) { allResults.push_back({}); } } } return allResults; } _logger.LogDebug("ANSTENSORRTPOSE::DetectObjectsBatch", "Processing batch of " + std::to_string(inputImages.size()) + " images", __FILE__, __LINE__); // Phase 1: Preprocess under brief lock BatchMetadata metadata; std::vector> inputs; { std::lock_guard lock(_mutex); inputs = PreprocessBatch(inputImages, metadata); } if (inputs.empty() || inputs[0].empty()) { _logger.LogFatal("ANSTENSORRTPOSE::DetectObjectsBatch", "Preprocessing failed", __FILE__, __LINE__); return {}; } // Phase 2: Inference -- mutex released; pool dispatches to idle GPU slot std::vector>> featureVectors; auto succ = m_trtEngine->runInference(inputs, featureVectors); if (!succ) { _logger.LogError("ANSTENSORRTPOSE::DetectObjectsBatch", "Error running inference", __FILE__, __LINE__); return {}; } // Phase 3: Parallel postprocessing -- each image is independent const size_t numBatch = featureVectors.size(); std::vector> batchDetections(numBatch); std::vector>> postFutures; postFutures.reserve(numBatch); for (size_t batchIdx = 0; batchIdx < numBatch; ++batchIdx) { const auto& batchOutput = featureVectors[batchIdx]; std::vector fv = batchOutput.empty() ? std::vector{} : batchOutput[0]; postFutures.push_back(std::async(std::launch::async, [this, fv = std::move(fv), cid = camera_id, idx = batchIdx, &metadata]() mutable { return PostProcessPoseBatch(fv, cid, idx, metadata); })); } for (size_t i = 0; i < numBatch; ++i) batchDetections[i] = postFutures[i].get(); if (_trackerEnabled) { for (auto& dets : batchDetections) { if (!dets.empty()) { dets = ApplyTracking(dets, camera_id); if (_stabilizationEnabled) dets = StabilizeDetections(dets, camera_id); } } } _logger.LogDebug("ANSTENSORRTPOSE::DetectObjectsBatch", "Batch processing complete. Images: " + std::to_string(numBatch), __FILE__, __LINE__); return batchDetections; } std::vector> ANSTENSORRTPOSE::PreprocessBatch(const std::vector& inputImages, BatchMetadata& outMetadata) { try { if (!_licenseValid) { _logger.LogFatal("ANSTENSORRTPOSE::PreprocessBatch", "Invalid license", __FILE__, __LINE__); return {}; } const auto& inputDims = m_trtEngine->getInputDims(); const int inputH = inputDims[0].d[1]; const int inputW = inputDims[0].d[2]; outMetadata.imgHeights.resize(inputImages.size()); outMetadata.imgWidths.resize(inputImages.size()); outMetadata.ratios.resize(inputImages.size()); std::vector batchProcessed; batchProcessed.reserve(inputImages.size()); cv::cuda::Stream stream; for (size_t i = 0; i < inputImages.size(); ++i) { const auto& inputImage = inputImages[i]; if (inputImage.empty()) { _logger.LogFatal("ANSTENSORRTPOSE::PreprocessBatch", "Empty input image at index " + std::to_string(i), __FILE__, __LINE__); return {}; } cv::cuda::GpuMat img; if (inputImage.channels() == 1) { cv::Mat img3Channel; cv::cvtColor(inputImage, img3Channel, cv::COLOR_GRAY2BGR); img.upload(img3Channel, stream); } else { img.upload(inputImage, stream); } cv::cuda::GpuMat imgRGB; cv::cuda::cvtColor(img, imgRGB, cv::COLOR_BGR2RGB, 0, stream); outMetadata.imgHeights[i] = imgRGB.rows; outMetadata.imgWidths[i] = imgRGB.cols; if (outMetadata.imgHeights[i] <= 0 || outMetadata.imgWidths[i] <= 0) { _logger.LogFatal("ANSTENSORRTPOSE::PreprocessBatch", "Image " + std::to_string(i) + " has invalid dimensions (Width: " + std::to_string(outMetadata.imgWidths[i]) + ", Height: " + std::to_string(outMetadata.imgHeights[i]) + ")", __FILE__, __LINE__); return {}; } outMetadata.ratios[i] = 1.f / std::min(inputW / static_cast(imgRGB.cols), inputH / static_cast(imgRGB.rows)); cv::cuda::GpuMat resized = imgRGB; if (resized.rows != inputH || resized.cols != inputW) { resized = Engine::resizeKeepAspectRatioPadRightBottom(imgRGB, inputH, inputW); } batchProcessed.push_back(std::move(resized)); } stream.waitForCompletion(); std::vector> inputs; inputs.push_back(std::move(batchProcessed)); return inputs; } catch (const std::exception& e) { _logger.LogFatal("ANSTENSORRTPOSE::PreprocessBatch", e.what(), __FILE__, __LINE__); return {}; } } std::vector ANSTENSORRTPOSE::PostProcessPoseBatch(std::vector& featureVector, const std::string& camera_id, size_t batchIdx, const BatchMetadata& metadata) { try { const auto& outputDims = m_trtEngine->getOutputDims(); auto numChannels = outputDims[0].d[1]; auto numAnchors = outputDims[0].d[2]; float ratio = metadata.ratios[batchIdx]; float imgWidth = static_cast(metadata.imgWidths[batchIdx]); float imgHeight = static_cast(metadata.imgHeights[batchIdx]); std::vector bboxes; std::vector scores; std::vector labels; std::vector indices; std::vector> kpss; cv::Mat output = cv::Mat(numChannels, numAnchors, CV_32F, featureVector.data()); output = output.t(); // Get all the YOLO proposals for (int i = 0; i < numAnchors; i++) { auto rowPtr = output.row(i).ptr(); auto bboxesPtr = rowPtr; auto scoresPtr = rowPtr + 4; auto kps_ptr = rowPtr + 5; float score = *scoresPtr; if (score > _modelConfig.detectionScoreThreshold) { float x = *bboxesPtr++; float y = *bboxesPtr++; float w = *bboxesPtr++; float h = *bboxesPtr; // Use batch-specific ratio and dimensions float x0 = std::clamp((x - 0.5f * w) * ratio, 0.f, imgWidth); float y0 = std::clamp((y - 0.5f * h) * ratio, 0.f, imgHeight); float x1 = std::clamp((x + 0.5f * w) * ratio, 0.f, imgWidth); float y1 = std::clamp((y + 0.5f * h) * ratio, 0.f, imgHeight); cv::Rect_ bbox; bbox.x = x0; bbox.y = y0; bbox.width = x1 - x0; bbox.height = y1 - y0; // Clamp to image boundaries bbox.x = std::max(0.f, bbox.x); bbox.y = std::max(0.f, bbox.y); bbox.width = std::min(imgWidth - bbox.x, bbox.width); bbox.height = std::min(imgHeight - bbox.y, bbox.height); // Process keypoints with batch-specific ratio and dimensions std::vector kps; for (int k = 0; k < NUM_KPS; k++) { float kpsX = *(kps_ptr + 3 * k) * ratio; float kpsY = *(kps_ptr + 3 * k + 1) * ratio; float kpsS = *(kps_ptr + 3 * k + 2); kpsX = std::clamp(kpsX, 0.f, imgWidth); kpsY = std::clamp(kpsY, 0.f, imgHeight); cv::Point2f kp(kpsX, kpsY); kps.push_back(kp); } bboxes.push_back(bbox); labels.push_back(0); // All detected objects are people scores.push_back(score); kpss.push_back(kps); } } // Run NMS cv::dnn::NMSBoxesBatched(bboxes, scores, labels, PROBABILITY_THRESHOLD, NMS_THRESHOLD, indices); std::vector objects; int classNameSize = static_cast(_classes.size()); // Choose the top k detections for (auto& chosenIdx : indices) { if (scores[chosenIdx] > _modelConfig.detectionScoreThreshold) { std::stringstream keypointXss; std::stringstream keypointYss; std::vector keypoints; for (size_t keypointIdx = 0; keypointIdx < kpss[chosenIdx].size(); keypointIdx++) { keypointXss << kpss[chosenIdx][keypointIdx].x; keypointYss << kpss[chosenIdx][keypointIdx].y; // Add semicolon after each value except the last one if (keypointIdx < kpss[chosenIdx].size() - 1) { keypointXss << ";"; keypointYss << ";"; } keypoints.push_back(kpss[chosenIdx][keypointIdx].x); keypoints.push_back(kpss[chosenIdx][keypointIdx].y); } std::string keypointXString = keypointXss.str(); std::string keypointYString = keypointYss.str(); std::string keypointString = keypointXString + "|" + keypointYString; Object obj{}; obj.confidence = scores[chosenIdx]; obj.classId = labels[chosenIdx]; if (!_classes.empty()) { if (obj.classId < classNameSize) { obj.className = _classes[obj.classId]; } else { obj.className = _classes[classNameSize - 1]; } } else { obj.className = "Unknown"; } obj.box = bboxes[chosenIdx]; obj.polygon = kpss[chosenIdx]; obj.kps = keypoints; obj.cameraId = camera_id; obj.extraInfo = keypointString; objects.push_back(obj); } } return objects; } catch (std::exception& e) { _logger.LogFatal("ANSTENSORRTPOSE::PostProcessPoseBatch", e.what(), __FILE__, __LINE__); return {}; } } std::vector> ANSTENSORRTPOSE::RunInferencesBatch( const std::vector& inputs, const std::string& camera_id) { { std::lock_guard lock(_mutex); if (!_modelLoadValid) { _logger.LogError("ANSTENSORRTPOSE::RunInferencesBatch", "Model not loaded", __FILE__, __LINE__); return {}; } if (!_licenseValid) { _logger.LogError("ANSTENSORRTPOSE::RunInferencesBatch", "Invalid license", __FILE__, __LINE__); return {}; } if (!_isInitialized) { _logger.LogError("ANSTENSORRTPOSE::RunInferencesBatch", "Engine not initialized", __FILE__, __LINE__); return {}; } if (inputs.empty()) return {}; } try { return DetectObjectsBatch(inputs, camera_id); } catch (const std::exception& e) { _logger.LogFatal("ANSTENSORRTPOSE::RunInferencesBatch", e.what(), __FILE__, __LINE__); return {}; } } } /*std::vector ANSTENSORRTPOSE::RunInference(const cv::Mat& inputImgBGR, const std::string& camera_id) { std::lock_guard lock(_mutex); if (!_modelLoadValid) { this->_logger.LogFatal("ANSTENSORRTPOSE::RunInference", "Cannot load the TensorRT model. Please check if it is exist", __FILE__, __LINE__); std::vector result; result.clear(); return result; } if (!_licenseValid) { this->_logger.LogFatal("ANSTENSORRTPOSE::RunInference", "Runtime license is not valid or expired. Please contact ANSCENTER", __FILE__, __LINE__); std::vector result; result.clear(); return result; } if (!_isInitialized) { this->_logger.LogFatal("ANSTENSORRTPOSE::RunInference", "Model is not initialized", __FILE__, __LINE__); std::vector result; result.clear(); return result; } try { std::vector result; if (inputImgBGR.empty()) return result; if ((inputImgBGR.cols < 10) || (inputImgBGR.rows < 10)) return result; return DetectObjects(inputImgBGR, camera_id); } catch (std::exception& e) { this->_logger.LogFatal("ANSTENSORRTPOSE::RunInference", e.what(), __FILE__, __LINE__); std::vector result; result.clear(); return result; } }*/