1578 lines
55 KiB
C++
1578 lines
55 KiB
C++
#include "ANSONNXOBB.h"
|
||
#include "EPLoader.h"
|
||
namespace ANSCENTER {
|
||
std::atomic<int> ANSONNXOBB::instanceCounter_(0); // Initialize static member
|
||
|
||
size_t ANSONNXOBB::vectorProduct(const std::vector<int64_t>& vector) {
|
||
return std::accumulate(vector.begin(), vector.end(), 1ull, std::multiplies<size_t>());
|
||
}
|
||
void ANSONNXOBB::letterBox(const cv::Mat& image, cv::Mat& outImage,
|
||
const cv::Size& newShape,
|
||
const cv::Scalar& color,
|
||
bool auto_,
|
||
bool scaleFill,
|
||
bool scaleUp,
|
||
int stride)
|
||
{
|
||
// Calculate the scaling ratio to fit the image within the new shape
|
||
float ratio = std::min(static_cast<float>(newShape.height) / image.rows,
|
||
static_cast<float>(newShape.width) / image.cols);
|
||
|
||
// Prevent scaling up if not allowed
|
||
if (!scaleUp) {
|
||
ratio = std::min(ratio, 1.0f);
|
||
}
|
||
|
||
// Calculate new dimensions after scaling
|
||
int newUnpadW = static_cast<int>(std::round(image.cols * ratio));
|
||
int newUnpadH = static_cast<int>(std::round(image.rows * ratio));
|
||
|
||
// Calculate padding needed to reach the desired shape
|
||
int dw = newShape.width - newUnpadW;
|
||
int dh = newShape.height - newUnpadH;
|
||
|
||
if (auto_) {
|
||
// Ensure padding is a multiple of stride for model compatibility
|
||
dw = (dw % stride) / 2;
|
||
dh = (dh % stride) / 2;
|
||
}
|
||
else if (scaleFill) {
|
||
// Scale to fill without maintaining aspect ratio
|
||
newUnpadW = newShape.width;
|
||
newUnpadH = newShape.height;
|
||
ratio = std::min(static_cast<float>(newShape.width) / image.cols,
|
||
static_cast<float>(newShape.height) / image.rows);
|
||
dw = 0;
|
||
dh = 0;
|
||
}
|
||
else {
|
||
// Evenly distribute padding on both sides
|
||
// Calculate separate padding for left/right and top/bottom to handle odd padding
|
||
int padLeft = dw / 2;
|
||
int padRight = dw - padLeft;
|
||
int padTop = dh / 2;
|
||
int padBottom = dh - padTop;
|
||
|
||
// Resize the image if the new dimensions differ
|
||
if (image.cols != newUnpadW || image.rows != newUnpadH) {
|
||
cv::resize(image, outImage, cv::Size(newUnpadW, newUnpadH), 0, 0, cv::INTER_LINEAR);
|
||
}
|
||
else {
|
||
// Avoid unnecessary copying if dimensions are the same
|
||
outImage = image;
|
||
}
|
||
|
||
// Apply padding to reach the desired shape
|
||
cv::copyMakeBorder(outImage, outImage, padTop, padBottom, padLeft, padRight, cv::BORDER_CONSTANT, color);
|
||
return; // Exit early since padding is already applied
|
||
}
|
||
|
||
// Resize the image if the new dimensions differ
|
||
if (image.cols != newUnpadW || image.rows != newUnpadH) {
|
||
cv::resize(image, outImage, cv::Size(newUnpadW, newUnpadH), 0, 0, cv::INTER_LINEAR);
|
||
}
|
||
else {
|
||
// Avoid unnecessary copying if dimensions are the same
|
||
outImage = image;
|
||
}
|
||
|
||
// Calculate separate padding for left/right and top/bottom to handle odd padding
|
||
int padLeft = dw / 2;
|
||
int padRight = dw - padLeft;
|
||
int padTop = dh / 2;
|
||
int padBottom = dh - padTop;
|
||
|
||
// Apply padding to reach the desired shape
|
||
cv::copyMakeBorder(outImage, outImage, padTop, padBottom, padLeft, padRight, cv::BORDER_CONSTANT, color);
|
||
}
|
||
void ANSONNXOBB::NMSBoxes(const std::vector<BoundingBox>& boundingBoxes,
|
||
const std::vector<float>& scores,
|
||
float scoreThreshold,
|
||
float nmsThreshold,
|
||
std::vector<int>& indices)
|
||
{
|
||
indices.clear();
|
||
|
||
const size_t numBoxes = boundingBoxes.size();
|
||
if (numBoxes == 0) {
|
||
DEBUG_PRINT("No bounding boxes to process in NMS");
|
||
return;
|
||
}
|
||
|
||
// Step 1: Filter out boxes with scores below the threshold
|
||
// and create a list of indices sorted by descending scores
|
||
std::vector<int> sortedIndices;
|
||
sortedIndices.reserve(numBoxes);
|
||
for (size_t i = 0; i < numBoxes; ++i) {
|
||
if (scores[i] >= scoreThreshold) {
|
||
sortedIndices.push_back(static_cast<int>(i));
|
||
}
|
||
}
|
||
|
||
// If no boxes remain after thresholding
|
||
if (sortedIndices.empty()) {
|
||
DEBUG_PRINT("No bounding boxes above score threshold");
|
||
return;
|
||
}
|
||
|
||
// Sort the indices based on scores in descending order
|
||
std::sort(sortedIndices.begin(), sortedIndices.end(),
|
||
[&scores](int idx1, int idx2) {
|
||
return scores[idx1] > scores[idx2];
|
||
});
|
||
|
||
// Step 2: Precompute the areas of all boxes
|
||
std::vector<float> areas(numBoxes, 0.0f);
|
||
for (size_t i = 0; i < numBoxes; ++i) {
|
||
areas[i] = boundingBoxes[i].width * boundingBoxes[i].height;
|
||
}
|
||
|
||
// Step 3: Suppression mask to mark boxes that are suppressed
|
||
std::vector<bool> suppressed(numBoxes, false);
|
||
|
||
// Step 4: Iterate through the sorted list and suppress boxes with high IoU
|
||
for (size_t i = 0; i < sortedIndices.size(); ++i) {
|
||
int currentIdx = sortedIndices[i];
|
||
if (suppressed[currentIdx]) {
|
||
continue;
|
||
}
|
||
|
||
// Select the current box as a valid detection
|
||
indices.push_back(currentIdx);
|
||
|
||
const BoundingBox& currentBox = boundingBoxes[currentIdx];
|
||
const float x1_max = currentBox.x;
|
||
const float y1_max = currentBox.y;
|
||
const float x2_max = currentBox.x + currentBox.width;
|
||
const float y2_max = currentBox.y + currentBox.height;
|
||
const float area_current = areas[currentIdx];
|
||
|
||
// Compare IoU of the current box with the rest
|
||
for (size_t j = i + 1; j < sortedIndices.size(); ++j) {
|
||
int compareIdx = sortedIndices[j];
|
||
if (suppressed[compareIdx]) {
|
||
continue;
|
||
}
|
||
|
||
const BoundingBox& compareBox = boundingBoxes[compareIdx];
|
||
const float x1 = std::max(x1_max, static_cast<float>(compareBox.x));
|
||
const float y1 = std::max(y1_max, static_cast<float>(compareBox.y));
|
||
const float x2 = std::min(x2_max, static_cast<float>(compareBox.x + compareBox.width));
|
||
const float y2 = std::min(y2_max, static_cast<float>(compareBox.y + compareBox.height));
|
||
|
||
const float interWidth = x2 - x1;
|
||
const float interHeight = y2 - y1;
|
||
|
||
if (interWidth <= 0 || interHeight <= 0) {
|
||
continue;
|
||
}
|
||
|
||
const float intersection = interWidth * interHeight;
|
||
const float unionArea = area_current + areas[compareIdx] - intersection;
|
||
const float iou = (unionArea > 0.0f) ? (intersection / unionArea) : 0.0f;
|
||
|
||
if (iou > nmsThreshold) {
|
||
suppressed[compareIdx] = true;
|
||
}
|
||
}
|
||
}
|
||
|
||
DEBUG_PRINT("NMS completed with " + std::to_string(indices.size()) + " indices remaining");
|
||
}
|
||
std::vector<cv::Point2f> ANSONNXOBB::OBBToPoints(const OrientedBoundingBox& obb) {
|
||
// Convert angle from radians to degrees for OpenCV
|
||
const float angleDeg = obb.angle * 180.0f / CV_PI;
|
||
|
||
// Create rotated rectangle from OBB parameters
|
||
const cv::RotatedRect rotatedRect(
|
||
cv::Point2f(obb.x, obb.y),
|
||
cv::Size2f(obb.width, obb.height),
|
||
angleDeg
|
||
);
|
||
|
||
// Extract corner points directly
|
||
std::vector<cv::Point2f> corners(4);
|
||
rotatedRect.points(corners.data());
|
||
|
||
return corners;
|
||
}
|
||
|
||
void ANSONNXOBB::drawBoundingBox(cv::Mat& image, const std::vector<Detection>& detections,
|
||
const std::vector<std::string>& classNames, const std::vector<cv::Scalar>& colors) {
|
||
for (const auto& detection : detections) {
|
||
if (detection.conf < _modelConfig.detectionScoreThreshold) continue;
|
||
if (detection.classId < 0 || static_cast<size_t>(detection.classId) >= classNames.size()) continue;
|
||
|
||
// Convert angle from radians to degrees for OpenCV
|
||
float angle_deg = detection.box.angle * 180.0f / CV_PI;
|
||
|
||
cv::RotatedRect rect(cv::Point2f(detection.box.x, detection.box.y),
|
||
cv::Size2f(detection.box.width, detection.box.height),
|
||
angle_deg);
|
||
|
||
// Convert rotated rectangle to polygon points
|
||
cv::Mat points_mat;
|
||
cv::boxPoints(rect, points_mat);
|
||
points_mat.convertTo(points_mat, CV_32SC1);
|
||
|
||
// Draw bounding box
|
||
cv::Scalar color = colors[detection.classId % colors.size()];
|
||
cv::polylines(image, points_mat, true, color, 3, cv::LINE_AA);
|
||
|
||
// Prepare label
|
||
std::string label = classNames[detection.classId] + ": " + cv::format("%.1f%%", detection.conf * 100);
|
||
int baseline = 0;
|
||
float fontScale = 0.6;
|
||
int thickness = 1;
|
||
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX, fontScale, thickness, &baseline);
|
||
|
||
// Calculate label position using bounding rect of rotated rectangle
|
||
cv::Rect brect = rect.boundingRect();
|
||
int x = brect.x;
|
||
int y = brect.y - labelSize.height - baseline;
|
||
|
||
// Adjust label position if it goes off-screen
|
||
if (y < 0) {
|
||
y = brect.y + brect.height;
|
||
if (y + labelSize.height > image.rows) {
|
||
y = image.rows - labelSize.height;
|
||
}
|
||
}
|
||
|
||
x = std::max(0, std::min(x, image.cols - labelSize.width));
|
||
|
||
// Draw label background (darker version of box color)
|
||
cv::Scalar labelBgColor = color * 0.6;
|
||
cv::rectangle(image, cv::Rect(x, y, labelSize.width, labelSize.height + baseline),
|
||
labelBgColor, cv::FILLED);
|
||
|
||
// Draw label text
|
||
cv::putText(image, label, cv::Point(x, y + labelSize.height),
|
||
cv::FONT_HERSHEY_DUPLEX, fontScale, cv::Scalar::all(255),
|
||
thickness, cv::LINE_AA);
|
||
}
|
||
}
|
||
std::vector<cv::Scalar> ANSONNXOBB::generateColors(const std::vector<std::string>& classNames, int seed) {
|
||
// Static cache to store colors based on class names to avoid regenerating
|
||
static std::unordered_map<size_t, std::vector<cv::Scalar>> colorCache;
|
||
|
||
// Compute a hash key based on class names to identify unique class configurations
|
||
size_t hashKey = 0;
|
||
for (const auto& name : classNames) {
|
||
hashKey ^= std::hash<std::string>{}(name)+0x9e3779b9 + (hashKey << 6) + (hashKey >> 2);
|
||
}
|
||
|
||
// Check if colors for this class configuration are already cached
|
||
auto it = colorCache.find(hashKey);
|
||
if (it != colorCache.end()) {
|
||
return it->second;
|
||
}
|
||
|
||
// Generate unique random colors for each class
|
||
std::vector<cv::Scalar> colors;
|
||
colors.reserve(classNames.size());
|
||
|
||
std::mt19937 rng(seed); // Initialize random number generator with fixed seed
|
||
std::uniform_int_distribution<int> uni(0, 255); // Define distribution for color values
|
||
|
||
for (size_t i = 0; i < classNames.size(); ++i) {
|
||
colors.emplace_back(cv::Scalar(uni(rng), uni(rng), uni(rng))); // Generate random BGR color
|
||
}
|
||
|
||
// Cache the generated colors for future use
|
||
colorCache.emplace(hashKey, colors);
|
||
|
||
return colorCache[hashKey];
|
||
}
|
||
|
||
void ANSONNXOBB::getCovarianceComponents(
|
||
const OrientedBoundingBox& box,
|
||
float& out1,
|
||
float& out2,
|
||
float& out3)
|
||
{
|
||
try {
|
||
// Validate input dimensions
|
||
if (box.width <= 0.0f || box.height <= 0.0f) {
|
||
this->_logger.LogError("ANSONNXOBB::getCovarianceComponents",
|
||
"[Instance " + std::to_string(instanceId_) + "] Invalid box dimensions: " +
|
||
std::to_string(box.width) + "x" + std::to_string(box.height),
|
||
__FILE__, __LINE__);
|
||
out1 = out2 = out3 = 0.0f;
|
||
return;
|
||
}
|
||
|
||
// Compute variance components (assuming uniform distribution in rectangle)
|
||
// For a rectangle with width w and height h:
|
||
// Variance along width axis: w²/12
|
||
// Variance along height axis: h²/12
|
||
const float varianceWidth = (box.width * box.width) / 12.0f;
|
||
const float varianceHeight = (box.height * box.height) / 12.0f;
|
||
|
||
// Precompute trigonometric values
|
||
const float cosTheta = std::cos(box.angle);
|
||
const float sinTheta = std::sin(box.angle);
|
||
const float cosSq = cosTheta * cosTheta;
|
||
const float sinSq = sinTheta * sinTheta;
|
||
const float sinCos = sinTheta * cosTheta;
|
||
|
||
// Compute rotated covariance matrix components
|
||
// After rotation by angle θ:
|
||
// σ_xx = a·cos²(θ) + b·sin²(θ)
|
||
// σ_yy = a·sin²(θ) + b·cos²(θ)
|
||
// σ_xy = (a - b)·sin(θ)·cos(θ)
|
||
out1 = varianceWidth * cosSq + varianceHeight * sinSq; // σ_xx
|
||
out2 = varianceWidth * sinSq + varianceHeight * cosSq; // σ_yy
|
||
out3 = (varianceWidth - varianceHeight) * sinCos; // σ_xy
|
||
|
||
// Validate outputs (check for NaN/Inf)
|
||
if (std::isnan(out1) || std::isnan(out2) || std::isnan(out3) ||
|
||
std::isinf(out1) || std::isinf(out2) || std::isinf(out3)) {
|
||
this->_logger.LogError("ANSONNXOBB::getCovarianceComponents",
|
||
"[Instance " + std::to_string(instanceId_) + "] Invalid output values (NaN/Inf)",
|
||
__FILE__, __LINE__);
|
||
out1 = out2 = out3 = 0.0f;
|
||
return;
|
||
}
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::getCovarianceComponents",
|
||
"[Instance " + std::to_string(instanceId_) + "] Exception: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
out1 = out2 = out3 = 0.0f;
|
||
}
|
||
}
|
||
std::vector<std::vector<float>> ANSONNXOBB::batchProbiou(
|
||
const std::vector<OrientedBoundingBox>& obb1,
|
||
const std::vector<OrientedBoundingBox>& obb2,
|
||
float eps)
|
||
{
|
||
try {
|
||
// Validate inputs
|
||
if (obb1.empty() || obb2.empty()) {
|
||
return {};
|
||
}
|
||
|
||
const size_t numBoxes1 = obb1.size();
|
||
const size_t numBoxes2 = obb2.size();
|
||
|
||
// Pre-allocate result matrix
|
||
std::vector<std::vector<float>> iouMatrix(numBoxes1, std::vector<float>(numBoxes2, 0.0f));
|
||
|
||
// Pre-compute covariance components for all boxes in obb1
|
||
std::vector<std::array<float, 5>> covData1(numBoxes1);
|
||
for (size_t i = 0; i < numBoxes1; ++i) {
|
||
const OrientedBoundingBox& box = obb1[i];
|
||
float a, b, c;
|
||
getCovarianceComponents(box, a, b, c);
|
||
covData1[i] = { box.x, box.y, a, b, c };
|
||
}
|
||
|
||
// Compute pairwise Prob-IoU
|
||
for (size_t i = 0; i < numBoxes1; ++i) {
|
||
const float x1 = covData1[i][0];
|
||
const float y1 = covData1[i][1];
|
||
const float a1 = covData1[i][2];
|
||
const float b1 = covData1[i][3];
|
||
const float c1 = covData1[i][4];
|
||
|
||
for (size_t j = 0; j < numBoxes2; ++j) {
|
||
const OrientedBoundingBox& box2 = obb2[j];
|
||
float a2, b2, c2;
|
||
getCovarianceComponents(box2, a2, b2, c2);
|
||
|
||
// Compute Bhattacharyya distance components
|
||
const float dx = x1 - box2.x;
|
||
const float dy = y1 - box2.y;
|
||
|
||
// Sum of covariance components
|
||
const float sumA = a1 + a2;
|
||
const float sumB = b1 + b2;
|
||
const float sumC = c1 + c2;
|
||
|
||
// Denominator for distance terms
|
||
const float denom = sumA * sumB - sumC * sumC + eps;
|
||
|
||
if (denom <= eps) {
|
||
// Degenerate covariance matrix, set IoU to 0
|
||
iouMatrix[i][j] = 0.0f;
|
||
continue;
|
||
}
|
||
|
||
// Mahalanobis distance term (T1)
|
||
const float t1 = ((sumA * dy * dy + sumB * dx * dx) * 0.25f) / denom;
|
||
|
||
// Cross term (T2)
|
||
const float t2 = ((sumC * dx * dy) * -0.5f) / denom;
|
||
|
||
// Determinant ratio term (T3)
|
||
const float det1 = a1 * b1 - c1 * c1;
|
||
const float det2 = a2 * b2 - c2 * c2;
|
||
|
||
// Ensure determinants are non-negative (numerical stability)
|
||
const float det1_safe = std::max(det1, 0.0f);
|
||
const float det2_safe = std::max(det2, 0.0f);
|
||
|
||
const float sqrtDetProduct = std::sqrt(det1_safe * det2_safe + eps);
|
||
const float numerator = sumA * sumB - sumC * sumC;
|
||
const float t3 = 0.5f * std::log((numerator / (4.0f * sqrtDetProduct)) + eps);
|
||
|
||
// Bhattacharyya distance
|
||
float bd = t1 + t2 + t3;
|
||
bd = std::clamp(bd, eps, 100.0f);
|
||
|
||
// Hellinger distance from Bhattacharyya distance
|
||
const float hd = std::sqrt(1.0f - std::exp(-bd) + eps);
|
||
|
||
// Convert Hellinger distance to IoU-like metric
|
||
iouMatrix[i][j] = 1.0f - hd;
|
||
}
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Computed "
|
||
<< numBoxes1 << "x" << numBoxes2 << " Prob-IoU matrix");
|
||
|
||
return iouMatrix;
|
||
}
|
||
catch (const std::bad_alloc& e) {
|
||
this->_logger.LogError("ANSONNXOBB::batchProbiou",
|
||
"[Instance " + std::to_string(instanceId_) + "] Memory allocation failed: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::batchProbiou",
|
||
"[Instance " + std::to_string(instanceId_) + "] Exception: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
std::vector<int> ANSONNXOBB::nmsRotatedImpl(
|
||
const std::vector<OrientedBoundingBox>& sortedBoxes,
|
||
float iouThreshold)
|
||
{
|
||
try {
|
||
if (sortedBoxes.empty()) {
|
||
return {};
|
||
}
|
||
|
||
const int numBoxes = static_cast<int>(sortedBoxes.size());
|
||
|
||
// Early return for single box
|
||
if (numBoxes == 1) {
|
||
return { 0 };
|
||
}
|
||
|
||
// Compute all pairwise IoU values
|
||
std::vector<std::vector<float>> iouMatrix = batchProbiou(sortedBoxes, sortedBoxes);
|
||
|
||
// Validate IoU matrix dimensions
|
||
if (iouMatrix.size() != static_cast<size_t>(numBoxes)) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotatedImpl",
|
||
"[Instance " + std::to_string(instanceId_) + "] IoU matrix size mismatch: " +
|
||
std::to_string(iouMatrix.size()) + " vs " + std::to_string(numBoxes),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
// Track which boxes to keep
|
||
std::vector<int> keepIndices;
|
||
keepIndices.reserve(numBoxes / 2); // Estimate ~50% will be kept
|
||
|
||
// NMS algorithm: keep boxes that don't overlap significantly with higher-scoring boxes
|
||
for (int j = 0; j < numBoxes; ++j) {
|
||
bool shouldKeep = true;
|
||
|
||
// Check against all previously kept boxes (higher scores due to sorting)
|
||
for (int i = 0; i < j; ++i) {
|
||
// Validate inner vector size
|
||
if (iouMatrix[i].size() != static_cast<size_t>(numBoxes)) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotatedImpl",
|
||
"[Instance " + std::to_string(instanceId_) + "] IoU matrix row " +
|
||
std::to_string(i) + " size mismatch",
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
// If current box overlaps significantly with a higher-scoring box, suppress it
|
||
if (iouMatrix[i][j] >= iouThreshold) {
|
||
shouldKeep = false;
|
||
break;
|
||
}
|
||
}
|
||
|
||
if (shouldKeep) {
|
||
keepIndices.push_back(j);
|
||
}
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] NMS kept "
|
||
<< keepIndices.size() << " of " << numBoxes
|
||
<< " boxes (threshold: " << iouThreshold << ")");
|
||
|
||
return keepIndices;
|
||
}
|
||
catch (const std::out_of_range& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotatedImpl",
|
||
"[Instance " + std::to_string(instanceId_) + "] Index out of range: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotatedImpl",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
std::vector<int> ANSONNXOBB::nmsRotated(
|
||
const std::vector<OrientedBoundingBox>& boxes,
|
||
const std::vector<float>& scores,
|
||
float iouThreshold)
|
||
{
|
||
try {
|
||
// Validate inputs
|
||
if (boxes.empty() || scores.empty()) {
|
||
return {};
|
||
}
|
||
|
||
if (boxes.size() != scores.size()) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotated",
|
||
"[Instance " + std::to_string(instanceId_) + "] Boxes and scores size mismatch: " +
|
||
std::to_string(boxes.size()) + " vs " + std::to_string(scores.size()),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
const size_t numBoxes = boxes.size();
|
||
|
||
// Create and sort indices by score (descending order)
|
||
std::vector<int> sortedIndices(numBoxes);
|
||
std::iota(sortedIndices.begin(), sortedIndices.end(), 0);
|
||
|
||
std::sort(sortedIndices.begin(), sortedIndices.end(),
|
||
[&scores](int a, int b) {
|
||
return scores[a] > scores[b];
|
||
});
|
||
|
||
// Create sorted boxes for NMS (avoid repeated indexing)
|
||
std::vector<OrientedBoundingBox> sortedBoxes;
|
||
sortedBoxes.reserve(numBoxes);
|
||
|
||
for (int idx : sortedIndices) {
|
||
sortedBoxes.push_back(boxes[idx]);
|
||
}
|
||
|
||
// Perform NMS on sorted boxes
|
||
std::vector<int> keepSorted = nmsRotatedImpl(sortedBoxes, iouThreshold);
|
||
|
||
// Map NMS results back to original indices
|
||
std::vector<int> keepOriginal;
|
||
keepOriginal.reserve(keepSorted.size());
|
||
|
||
for (int sortedIdx : keepSorted) {
|
||
keepOriginal.push_back(sortedIndices[sortedIdx]);
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] NMS kept "
|
||
<< keepOriginal.size() << " of " << numBoxes << " boxes");
|
||
|
||
return keepOriginal;
|
||
}
|
||
catch (const std::out_of_range& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotated",
|
||
"[Instance " + std::to_string(instanceId_) + "] Index out of range: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nmsRotated",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
std::vector<Object> ANSONNXOBB::nonMaxSuppression(
|
||
const std::vector<Detection>& inputDetections,
|
||
const std::string& camera_id,
|
||
float confThreshold,
|
||
float iouThreshold,
|
||
int maxDetections)
|
||
{
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
|
||
try {
|
||
if (inputDetections.empty()) {
|
||
return {};
|
||
}
|
||
|
||
// Filter by confidence threshold and pre-allocate
|
||
std::vector<Detection> candidates;
|
||
candidates.reserve(inputDetections.size());
|
||
|
||
for (const auto& det : inputDetections) {
|
||
if (det.conf > confThreshold) {
|
||
candidates.push_back(det);
|
||
}
|
||
}
|
||
|
||
if (candidates.empty()) {
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] No detections passed confidence threshold");
|
||
return {};
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] " << candidates.size()
|
||
<< " candidates for NMS (from " << inputDetections.size() << " total)");
|
||
|
||
// Extract boxes and scores for NMS
|
||
std::vector<OrientedBoundingBox> boxes;
|
||
std::vector<float> scores;
|
||
boxes.reserve(candidates.size());
|
||
scores.reserve(candidates.size());
|
||
|
||
for (const auto& det : candidates) {
|
||
boxes.push_back(det.box);
|
||
scores.push_back(det.conf);
|
||
}
|
||
|
||
// Perform rotated NMS
|
||
std::vector<int> keepIndices = nmsRotated(boxes, scores, iouThreshold);
|
||
|
||
// Limit to max detections
|
||
const size_t finalCount = std::min(static_cast<size_t>(maxDetections), keepIndices.size());
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] " << finalCount
|
||
<< " detections after NMS");
|
||
|
||
// Build final results
|
||
std::vector<Object> results;
|
||
results.reserve(finalCount);
|
||
|
||
for (size_t i = 0; i < finalCount; ++i) {
|
||
const int idx = keepIndices[i];
|
||
const Detection& det = candidates[idx];
|
||
const OrientedBoundingBox& obb = det.box;
|
||
|
||
Object obj;
|
||
obj.classId = det.classId;
|
||
obj.confidence = det.conf;
|
||
obj.cameraId = camera_id;
|
||
|
||
// Store OBB parameters as keypoints
|
||
obj.kps = { obb.x, obb.y, obb.width, obb.height, obb.angle };
|
||
|
||
// Convert OBB to polygon points
|
||
obj.polygon = OBBToPoints(obb);
|
||
|
||
// Calculate axis-aligned bounding box
|
||
obj.box = cv::boundingRect(obj.polygon);
|
||
|
||
results.push_back(std::move(obj));
|
||
}
|
||
|
||
return results;
|
||
}
|
||
catch (const cv::Exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nonMaxSuppression",
|
||
"[Instance " + std::to_string(instanceId_) + "] OpenCV error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::nonMaxSuppression",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
bool ANSONNXOBB::Init(const std::string& modelPath, bool useGPU, int deviceId)
|
||
{
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
try {
|
||
deviceId_ = deviceId;
|
||
|
||
const auto& ep = ANSCENTER::EPLoader::Current();
|
||
if (Ort::Global<void>::api_ == nullptr)
|
||
Ort::InitApi(static_cast<const OrtApi*>(EPLoader::GetOrtApiRaw()));
|
||
std::cout << "[ANSONNXOBB] EP ready: "
|
||
<< ANSCENTER::EPLoader::EngineTypeName(ep.type) << std::endl;
|
||
|
||
// Unique environment name per instance to avoid conflicts
|
||
std::string envName = "ONNX_OBB_INST" + std::to_string(instanceId_);
|
||
env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, envName.c_str());
|
||
|
||
sessionOptions = Ort::SessionOptions();
|
||
sessionOptions.SetIntraOpNumThreads(
|
||
std::min(6, static_cast<int>(std::thread::hardware_concurrency())));
|
||
sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||
|
||
// ── Log available providers ─────────────────────────────────────────
|
||
std::vector<std::string> availableProviders = Ort::GetAvailableProviders();
|
||
std::cout << "[Instance " << instanceId_ << "] Available Execution Providers:" << std::endl;
|
||
for (const auto& p : availableProviders)
|
||
std::cout << " - " << p << std::endl;
|
||
|
||
// ── Attach EP based on runtime-detected hardware ────────────────────
|
||
if (useGPU) {
|
||
bool attached = false;
|
||
|
||
switch (ep.type) {
|
||
|
||
case ANSCENTER::EngineType::NVIDIA_GPU: {
|
||
auto it = std::find(availableProviders.begin(),
|
||
availableProviders.end(), "CUDAExecutionProvider");
|
||
if (it == availableProviders.end()) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", "CUDAExecutionProvider not in DLL — "
|
||
"check ep/cuda/ has the CUDA ORT build.", __FILE__, __LINE__);
|
||
break;
|
||
}
|
||
try {
|
||
OrtCUDAProviderOptionsV2* cuda_options = nullptr;
|
||
Ort::GetApi().CreateCUDAProviderOptions(&cuda_options);
|
||
|
||
std::string deviceIdStr = std::to_string(deviceId_);
|
||
const char* keys[] = { "device_id" };
|
||
const char* values[] = { deviceIdStr.c_str() };
|
||
Ort::GetApi().UpdateCUDAProviderOptions(cuda_options, keys, values, 1);
|
||
|
||
sessionOptions.AppendExecutionProvider_CUDA_V2(*cuda_options);
|
||
Ort::GetApi().ReleaseCUDAProviderOptions(cuda_options);
|
||
|
||
std::cout << "[Instance " << instanceId_ << "] CUDA EP attached on device "
|
||
<< deviceId_ << "." << std::endl;
|
||
attached = true;
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", e.what(), __FILE__, __LINE__);
|
||
}
|
||
break;
|
||
}
|
||
|
||
case ANSCENTER::EngineType::AMD_GPU: {
|
||
auto it = std::find(availableProviders.begin(),
|
||
availableProviders.end(), "DmlExecutionProvider");
|
||
if (it == availableProviders.end()) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", "DmlExecutionProvider not in DLL — "
|
||
"check ep/directml/ has the DirectML ORT build.", __FILE__, __LINE__);
|
||
break;
|
||
}
|
||
try {
|
||
std::unordered_map<std::string, std::string> opts = {
|
||
{ "device_id", std::to_string(deviceId_) }
|
||
};
|
||
sessionOptions.AppendExecutionProvider("DML", opts);
|
||
std::cout << "[Instance " << instanceId_ << "] DirectML EP attached on device "
|
||
<< deviceId_ << "." << std::endl;
|
||
attached = true;
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", e.what(), __FILE__, __LINE__);
|
||
}
|
||
break;
|
||
}
|
||
|
||
case ANSCENTER::EngineType::OPENVINO_GPU: {
|
||
auto it = std::find(availableProviders.begin(),
|
||
availableProviders.end(), "OpenVINOExecutionProvider");
|
||
if (it == availableProviders.end()) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", "OpenVINOExecutionProvider not in DLL — "
|
||
"check ep/openvino/ has the OpenVINO ORT build.", __FILE__, __LINE__);
|
||
break;
|
||
}
|
||
|
||
// FP32 + single thread preserved for determinism; each instance gets its own stream and cache
|
||
const std::string precision = "FP32";
|
||
const std::string numberOfThreads = "1";
|
||
const std::string numberOfStreams = std::to_string(instanceId_ + 1);
|
||
const std::string primaryDevice = "GPU." + std::to_string(deviceId_);
|
||
const std::string cacheDir = "./ov_cache_inst" + std::to_string(instanceId_);
|
||
|
||
std::vector<std::unordered_map<std::string, std::string>> try_configs = {
|
||
{ {"device_type", primaryDevice}, {"precision",precision},
|
||
{"num_of_threads",numberOfThreads}, {"num_streams",numberOfStreams},
|
||
{"enable_opencl_throttling","False"}, {"enable_qdq_optimizer","False"},
|
||
{"cache_dir", cacheDir} },
|
||
{ {"device_type","GPU"}, {"precision",precision},
|
||
{"num_of_threads",numberOfThreads}, {"num_streams",numberOfStreams},
|
||
{"enable_opencl_throttling","False"}, {"enable_qdq_optimizer","False"},
|
||
{"cache_dir", cacheDir} },
|
||
{ {"device_type","AUTO:GPU,CPU"}, {"precision",precision},
|
||
{"num_of_threads",numberOfThreads}, {"num_streams",numberOfStreams},
|
||
{"enable_opencl_throttling","False"}, {"enable_qdq_optimizer","False"},
|
||
{"cache_dir", cacheDir} }
|
||
};
|
||
|
||
for (const auto& config : try_configs) {
|
||
try {
|
||
sessionOptions.AppendExecutionProvider_OpenVINO_V2(config);
|
||
std::cout << "[Instance " << instanceId_ << "] OpenVINO EP attached ("
|
||
<< config.at("device_type") << ", stream: " << numberOfStreams << ")." << std::endl;
|
||
attached = true;
|
||
break;
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
this->_logger.LogError("ANSONNXOBB::Init", e.what(), __FILE__, __LINE__);
|
||
}
|
||
}
|
||
|
||
if (!attached)
|
||
std::cerr << "[Instance " << instanceId_ << "] OpenVINO EP: all device configs failed." << std::endl;
|
||
break;
|
||
}
|
||
|
||
default:
|
||
break;
|
||
}
|
||
|
||
if (!attached) {
|
||
std::cerr << "[Instance " << instanceId_ << "] No GPU EP attached — running on CPU." << std::endl;
|
||
this->_logger.LogFatal("ANSONNXOBB::Init", "GPU EP not attached. Running on CPU.", __FILE__, __LINE__);
|
||
}
|
||
}
|
||
else {
|
||
std::cout << "[Instance " << instanceId_ << "] Inference device: CPU (useGPU=false)" << std::endl;
|
||
}
|
||
|
||
// ── Load model ──────────────────────────────────────────────────────
|
||
#ifdef _WIN32
|
||
std::wstring w_modelPath = std::wstring(modelPath.begin(), modelPath.end());
|
||
session = Ort::Session(env, w_modelPath.c_str(), sessionOptions);
|
||
#else
|
||
session = Ort::Session(env, modelPath.c_str(), sessionOptions);
|
||
#endif
|
||
|
||
Ort::AllocatorWithDefaultOptions allocator;
|
||
|
||
// ── Input shape ─────────────────────────────────────────────────────
|
||
Ort::TypeInfo inputTypeInfo = session.GetInputTypeInfo(0);
|
||
std::vector<int64_t> inputTensorShapeVec =
|
||
inputTypeInfo.GetTensorTypeAndShapeInfo().GetShape();
|
||
isDynamicInputShape = (inputTensorShapeVec.size() >= 4) &&
|
||
(inputTensorShapeVec[2] == -1 && inputTensorShapeVec[3] == -1);
|
||
|
||
// ── Node names ──────────────────────────────────────────────────────
|
||
auto input_name = session.GetInputNameAllocated(0, allocator);
|
||
inputNodeNameAllocatedStrings.push_back(std::move(input_name));
|
||
inputNames.push_back(inputNodeNameAllocatedStrings.back().get());
|
||
|
||
auto output_name = session.GetOutputNameAllocated(0, allocator);
|
||
outputNodeNameAllocatedStrings.push_back(std::move(output_name));
|
||
outputNames.push_back(outputNodeNameAllocatedStrings.back().get());
|
||
|
||
// ── Input image size ────────────────────────────────────────────────
|
||
if (inputTensorShapeVec.size() >= 4) {
|
||
inputImageShape = cv::Size(
|
||
static_cast<int>(inputTensorShapeVec[3]),
|
||
static_cast<int>(inputTensorShapeVec[2]));
|
||
}
|
||
else {
|
||
throw std::runtime_error("Invalid input tensor shape.");
|
||
}
|
||
|
||
numInputNodes = session.GetInputCount();
|
||
numOutputNodes = session.GetOutputCount();
|
||
|
||
classColors = generateColors(_classes);
|
||
|
||
std::cout << "[Instance " << instanceId_ << "] Model loaded successfully with "
|
||
<< numInputNodes << " input nodes and " << numOutputNodes << " output nodes." << std::endl;
|
||
|
||
// ── Warmup ──────────────────────────────────────────────────────────
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Starting warmup...");
|
||
warmupModel();
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Warmup completed successfully.");
|
||
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::Init",
|
||
std::string("[Instance ") + std::to_string(instanceId_) + "] " + e.what(),
|
||
__FILE__, __LINE__);
|
||
return false;
|
||
}
|
||
}
|
||
void ANSONNXOBB::warmupModel() {
|
||
try {
|
||
// Create dummy input image with correct size
|
||
cv::Mat dummyImage = cv::Mat::zeros(inputImageShape.height, inputImageShape.width, CV_8UC3);
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Warmup: dummy image "
|
||
<< dummyImage.cols << "x" << dummyImage.rows);
|
||
|
||
// Run 3 warmup inferences to stabilize
|
||
for (int i = 0; i < 3; ++i) {
|
||
try {
|
||
// Your preprocessing logic here
|
||
float* blob = nullptr;
|
||
std::vector<int64_t> inputShape;
|
||
|
||
// If you have a preprocess method, call it
|
||
// Otherwise, create a simple dummy tensor
|
||
size_t tensorSize = 1 * 3 * inputImageShape.height * inputImageShape.width;
|
||
blob = new float[tensorSize];
|
||
std::memset(blob, 0, tensorSize * sizeof(float));
|
||
|
||
inputShape = { 1, 3, inputImageShape.height, inputImageShape.width };
|
||
|
||
// Create input tensor
|
||
Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||
Ort::Value inputTensor = Ort::Value::CreateTensor<float>(
|
||
memoryInfo,
|
||
blob,
|
||
tensorSize,
|
||
inputShape.data(),
|
||
inputShape.size()
|
||
);
|
||
|
||
// Run inference
|
||
std::vector<Ort::Value> outputTensors = session.Run(
|
||
Ort::RunOptions{ nullptr },
|
||
inputNames.data(),
|
||
&inputTensor,
|
||
1,
|
||
outputNames.data(),
|
||
numOutputNodes
|
||
);
|
||
|
||
// Clean up
|
||
delete[] blob;
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Warmup " << (i + 1) << "/3 completed");
|
||
}
|
||
catch (const std::exception& e) {
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Warmup iteration " << i
|
||
<< " failed (non-critical): " << e.what());
|
||
}
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Warmup successful - all states initialized");
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogWarn("ANSONNXOBB::warmupModel",
|
||
std::string("[Instance ") + std::to_string(instanceId_) + "] Warmup failed: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
}
|
||
}
|
||
cv::Mat ANSONNXOBB::preprocess(const cv::Mat& image, float*& blob, std::vector<int64_t>& inputTensorShape) {
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
|
||
// Clean up existing blob first to prevent memory leak
|
||
if (blob != nullptr) {
|
||
delete[] blob;
|
||
blob = nullptr;
|
||
}
|
||
|
||
try {
|
||
// Validate input image
|
||
if (image.empty() || image.data == nullptr) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"Input image is empty or has null data pointer", __FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
if (image.cols <= 0 || image.rows <= 0) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"Invalid image dimensions: " + std::to_string(image.cols) + "x" + std::to_string(image.rows),
|
||
__FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
// Check for NaN/Inf values (only in debug builds for performance)
|
||
#ifdef DEBUG_MODE
|
||
double minVal, maxVal;
|
||
cv::minMaxLoc(image, &minVal, &maxVal);
|
||
|
||
if (std::isnan(minVal) || std::isnan(maxVal) || std::isinf(minVal) || std::isinf(maxVal)) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"Input image contains NaN or Inf values. Range: [" + std::to_string(minVal) +
|
||
", " + std::to_string(maxVal) + "]", __FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Input: " << image.cols << "x" << image.rows
|
||
<< ", channels: " << image.channels() << ", type: " << image.type()
|
||
<< ", range: [" << minVal << ", " << maxVal << "]");
|
||
#else
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Input: " << image.cols << "x" << image.rows
|
||
<< ", channels: " << image.channels());
|
||
#endif
|
||
|
||
// Resize and pad image using letterBox
|
||
cv::Mat resizedImage;
|
||
letterBox(image, resizedImage, inputImageShape, cv::Scalar(114, 114, 114),
|
||
isDynamicInputShape, false, true, 32);
|
||
|
||
if (resizedImage.empty()) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"letterBox returned empty image", __FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
// Update input tensor shape
|
||
inputTensorShape[2] = resizedImage.rows;
|
||
inputTensorShape[3] = resizedImage.cols;
|
||
|
||
// Validate resized dimensions
|
||
const int channels = resizedImage.channels();
|
||
const int height = resizedImage.rows;
|
||
const int width = resizedImage.cols;
|
||
|
||
if (channels != 3) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"Expected 3 channels, got " + std::to_string(channels), __FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
// Calculate memory requirements
|
||
const size_t imageSize = static_cast<size_t>(width) * height;
|
||
const size_t totalSize = imageSize * channels;
|
||
|
||
// Check for potential overflow
|
||
if (totalSize > SIZE_MAX / sizeof(float)) {
|
||
this->_logger.LogError("ANSONNXOBB::preprocess",
|
||
"Image size too large: would overflow memory allocation", __FILE__, __LINE__);
|
||
return cv::Mat();
|
||
}
|
||
|
||
// Allocate blob memory for CHW format
|
||
blob = new float[totalSize];
|
||
|
||
// Convert to float and normalize in one operation
|
||
resizedImage.convertTo(resizedImage, CV_32FC3, 1.0 / 255.0);
|
||
|
||
// Convert from HWC (OpenCV) to CHW (ONNX) format efficiently
|
||
std::vector<cv::Mat> channelMats(channels);
|
||
for (int c = 0; c < channels; ++c) {
|
||
channelMats[c] = cv::Mat(height, width, CV_32FC1, blob + c * imageSize);
|
||
}
|
||
cv::split(resizedImage, channelMats);
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Preprocessing completed: "
|
||
<< width << "x" << height);
|
||
|
||
return resizedImage;
|
||
}
|
||
catch (const cv::Exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::preprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] OpenCV error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
if (blob != nullptr) {
|
||
delete[] blob;
|
||
blob = nullptr;
|
||
}
|
||
return cv::Mat();
|
||
}
|
||
catch (const std::bad_alloc& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::preprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] Memory allocation failed for " +
|
||
std::to_string(static_cast<size_t>(inputTensorShape[2]) * inputTensorShape[3] * 3 * sizeof(float)) +
|
||
" bytes: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
if (blob != nullptr) {
|
||
delete[] blob;
|
||
blob = nullptr;
|
||
}
|
||
return cv::Mat();
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::preprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " + e.what(),
|
||
__FILE__, __LINE__);
|
||
if (blob != nullptr) {
|
||
delete[] blob;
|
||
blob = nullptr;
|
||
}
|
||
return cv::Mat();
|
||
}
|
||
}
|
||
|
||
std::vector<Object> ANSONNXOBB::postprocess(
|
||
const cv::Size& originalImageSize,
|
||
const cv::Size& resizedImageShape,
|
||
const std::vector<Ort::Value>& outputTensors,
|
||
int topk,
|
||
const std::string& camera_id)
|
||
{
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
|
||
try {
|
||
// Validate output tensors
|
||
if (outputTensors.empty()) {
|
||
this->_logger.LogError("ANSONNXOBB::postprocess",
|
||
"Output tensors are empty", __FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
// Extract output tensor data and shape [1, num_features, num_detections]
|
||
const float* rawOutput = outputTensors[0].GetTensorData<float>();
|
||
const std::vector<int64_t> outputShape = outputTensors[0].GetTensorTypeAndShapeInfo().GetShape();
|
||
|
||
if (outputShape.size() < 3) {
|
||
this->_logger.LogError("ANSONNXOBB::postprocess",
|
||
"Invalid output shape dimensions: " + std::to_string(outputShape.size()),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
const int numFeatures = static_cast<int>(outputShape[1]);
|
||
const int numDetections = static_cast<int>(outputShape[2]);
|
||
|
||
if (numDetections == 0) {
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] No detections in output");
|
||
return {};
|
||
}
|
||
|
||
// Calculate number of class labels (layout: [x, y, w, h, scores..., angle])
|
||
const int numLabels = numFeatures - 5;
|
||
if (numLabels <= 0) {
|
||
this->_logger.LogError("ANSONNXOBB::postprocess",
|
||
"Invalid number of labels: " + std::to_string(numLabels),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Processing " << numDetections
|
||
<< " detections with " << numLabels << " classes");
|
||
|
||
// Calculate letterbox transformation parameters
|
||
const float inputWidth = static_cast<float>(resizedImageShape.width);
|
||
const float inputHeight = static_cast<float>(resizedImageShape.height);
|
||
const float originalWidth = static_cast<float>(originalImageSize.width);
|
||
const float originalHeight = static_cast<float>(originalImageSize.height);
|
||
|
||
const float scale = std::min(inputHeight / originalHeight, inputWidth / originalWidth);
|
||
const float paddedWidth = std::round(originalWidth * scale);
|
||
const float paddedHeight = std::round(originalHeight * scale);
|
||
const float offsetX = (inputWidth - paddedWidth) / 2.0f;
|
||
const float offsetY = (inputHeight - paddedHeight) / 2.0f;
|
||
const float inverseScale = 1.0f / scale;
|
||
|
||
// Transpose output: [num_features, num_detections] -> [num_detections, num_features]
|
||
cv::Mat output = cv::Mat(numFeatures, numDetections, CV_32F, const_cast<float*>(rawOutput));
|
||
output = output.t();
|
||
|
||
// Pre-allocate vectors with reasonable capacity
|
||
std::vector<Detection> candidateDetections;
|
||
candidateDetections.reserve(numDetections / 2); // Estimate ~50% will pass threshold
|
||
|
||
// Extract and filter detections
|
||
for (int i = 0; i < numDetections; ++i) {
|
||
const float* rowPtr = output.ptr<float>(i);
|
||
|
||
// Extract bounding box parameters (in letterbox space)
|
||
const float x = rowPtr[0];
|
||
const float y = rowPtr[1];
|
||
const float w = rowPtr[2];
|
||
const float h = rowPtr[3];
|
||
|
||
// Find best class and score
|
||
const float* scoresPtr = rowPtr + 4;
|
||
float maxScore = -FLT_MAX;
|
||
int classId = -1;
|
||
|
||
for (int j = 0; j < numLabels; ++j) {
|
||
const float score = scoresPtr[j];
|
||
if (score > maxScore) {
|
||
maxScore = score;
|
||
classId = j;
|
||
}
|
||
}
|
||
|
||
// Apply score threshold
|
||
if (maxScore <= _modelConfig.detectionScoreThreshold) {
|
||
continue;
|
||
}
|
||
|
||
// Extract rotation angle (stored after class scores)
|
||
const float angle = rowPtr[4 + numLabels];
|
||
|
||
// Transform coordinates from letterbox space to original image space
|
||
const float cx = (x - offsetX) * inverseScale;
|
||
const float cy = (y - offsetY) * inverseScale;
|
||
const float bw = w * inverseScale;
|
||
const float bh = h * inverseScale;
|
||
|
||
// Clamp coordinates to image bounds
|
||
const float clampedCx = std::clamp(cx, 0.0f, originalWidth);
|
||
const float clampedCy = std::clamp(cy, 0.0f, originalHeight);
|
||
const float clampedWidth = std::clamp(bw, 0.0f, originalWidth);
|
||
const float clampedHeight = std::clamp(bh, 0.0f, originalHeight);
|
||
|
||
// Create detection
|
||
OrientedBoundingBox obb(clampedCx, clampedCy, clampedWidth, clampedHeight, angle);
|
||
candidateDetections.emplace_back(Detection{ obb, maxScore, classId });
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] " << candidateDetections.size()
|
||
<< " detections passed score threshold");
|
||
|
||
// Apply Non-Maximum Suppression
|
||
std::vector<Object> finalDetections = nonMaxSuppression(
|
||
candidateDetections,
|
||
camera_id,
|
||
_modelConfig.modelConfThreshold,
|
||
_modelConfig.modelMNSThreshold,
|
||
topk
|
||
);
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] " << finalDetections.size()
|
||
<< " detections after NMS");
|
||
|
||
return finalDetections;
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::postprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] ONNX Runtime error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const cv::Exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::postprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] OpenCV error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::postprocess",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
|
||
std::vector<Object> ANSONNXOBB::detect(const cv::Mat& image, const std::string& camera_id) {
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
float* blobPtr = nullptr;
|
||
|
||
try {
|
||
// Validate input image
|
||
if (image.empty() || image.data == nullptr) {
|
||
this->_logger.LogError("ANSONNXOBB::detect",
|
||
"Input image is empty or has null data pointer", __FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
if (image.cols <= 0 || image.rows <= 0) {
|
||
this->_logger.LogError("ANSONNXOBB::detect",
|
||
"Invalid image dimensions: " + std::to_string(image.cols) + "x" + std::to_string(image.rows),
|
||
__FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Detecting objects in "
|
||
<< image.cols << "x" << image.rows << " image");
|
||
|
||
// Prepare input tensor shape (batch size, channels, height, width)
|
||
std::vector<int64_t> inputTensorShape = { 1, 3, inputImageShape.height, inputImageShape.width };
|
||
|
||
// Preprocess image and get blob pointer
|
||
cv::Mat preprocessedImage = preprocess(image, blobPtr, inputTensorShape);
|
||
if (preprocessedImage.empty() || blobPtr == nullptr) {
|
||
this->_logger.LogError("ANSONNXOBB::detect",
|
||
"Preprocessing failed", __FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
|
||
// Calculate input tensor size
|
||
const size_t inputTensorSize = vectorProduct(inputTensorShape);
|
||
|
||
// Create ONNX Runtime memory info (static to avoid repeated allocation)
|
||
static Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(
|
||
OrtArenaAllocator, OrtMemTypeDefault);
|
||
|
||
// Create input tensor directly from blob pointer (avoid vector copy)
|
||
Ort::Value inputTensor = Ort::Value::CreateTensor<float>(
|
||
memoryInfo,
|
||
blobPtr,
|
||
inputTensorSize,
|
||
inputTensorShape.data(),
|
||
inputTensorShape.size()
|
||
);
|
||
|
||
// Run inference
|
||
std::vector<Ort::Value> outputTensors = session.Run(
|
||
Ort::RunOptions{ nullptr },
|
||
inputNames.data(),
|
||
&inputTensor,
|
||
numInputNodes,
|
||
outputNames.data(),
|
||
numOutputNodes
|
||
);
|
||
|
||
// Clean up blob after inference
|
||
delete[] blobPtr;
|
||
blobPtr = nullptr;
|
||
|
||
// Post-process results
|
||
const cv::Size resizedImageShape(
|
||
static_cast<int>(inputTensorShape[3]),
|
||
static_cast<int>(inputTensorShape[2])
|
||
);
|
||
|
||
std::vector<Object> detections = postprocess(
|
||
image.size(),
|
||
resizedImageShape,
|
||
outputTensors,
|
||
100,
|
||
camera_id
|
||
);
|
||
|
||
DEBUG_PRINT("[Instance " << instanceId_ << "] Detected "
|
||
<< detections.size() << " objects");
|
||
|
||
return detections;
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::detect",
|
||
"[Instance " + std::to_string(instanceId_) + "] ONNX Runtime error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
if (blobPtr != nullptr) {
|
||
delete[] blobPtr;
|
||
}
|
||
return {};
|
||
}
|
||
catch (const cv::Exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::detect",
|
||
"[Instance " + std::to_string(instanceId_) + "] OpenCV error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
if (blobPtr != nullptr) {
|
||
delete[] blobPtr;
|
||
}
|
||
return {};
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::detect",
|
||
"[Instance " + std::to_string(instanceId_) + "] Error: " +
|
||
std::string(e.what()),
|
||
__FILE__, __LINE__);
|
||
if (blobPtr != nullptr) {
|
||
delete[] blobPtr;
|
||
}
|
||
return {};
|
||
}
|
||
}
|
||
// Public functions
|
||
ANSONNXOBB::~ANSONNXOBB() {
|
||
Destroy();
|
||
}
|
||
bool ANSONNXOBB::Destroy() {
|
||
std::cout << "[ANSONNXOBB] Destroyed instance " << instanceId_ << std::endl;
|
||
return true;
|
||
}
|
||
bool ANSONNXOBB::OptimizeModel(bool fp16, std::string& optimizedModelFolder) {
|
||
if (!ANSODBase::OptimizeModel(fp16, optimizedModelFolder)) {
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
bool ANSONNXOBB::Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) {
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
ModelLoadingGuard mlg(_modelLoading);
|
||
try {
|
||
_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::ONNXPOSE;
|
||
_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 = (modelConfig.precisionType == PrecisionType::FP16);
|
||
|
||
if (FileExist(_modelConfigFile)) {
|
||
ModelType modelType;
|
||
std::vector<int> 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("ANSONNXCL::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromString();
|
||
}
|
||
else {
|
||
this->_logger.LogDebug("ANSONNXCL::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromFile();
|
||
}
|
||
}
|
||
// 1. Load labelMap and engine
|
||
labelMap.clear();
|
||
if (!_classes.empty())
|
||
labelMap = VectorToCommaSeparatedString(_classes);
|
||
|
||
// 2. Initialize ONNX Runtime session
|
||
instanceId_ = instanceCounter_.fetch_add(1); // Atomic increment
|
||
result = Init(_modelFilePath, true, 0);
|
||
_modelLoadValid = true;
|
||
_isInitialized = true;
|
||
return result;
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXCL::Initialize", e.what(), __FILE__, __LINE__);
|
||
return false;
|
||
}
|
||
}
|
||
bool ANSONNXOBB::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) {
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
ModelLoadingGuard mlg(_modelLoading);
|
||
try {
|
||
bool result = ANSODBase::LoadModel(modelZipFilePath, modelZipPassword);
|
||
if (!result) return false;
|
||
_modelConfig.detectionType = ANSCENTER::DetectionType::CLASSIFICATION;
|
||
_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
|
||
// if (_modelConfig.precisionType == PrecisionType::FP16)_fp16 = true;
|
||
_fp16 = true; // Load Model from Here
|
||
|
||
// 0. Check if the configuration file exist
|
||
if (FileExist(_modelConfigFile)) {
|
||
ModelType modelType;
|
||
std::vector<int> 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("ANSONNXOBB::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromString();
|
||
}
|
||
else {
|
||
this->_logger.LogDebug("ANSONNXOBB::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromFile();
|
||
}
|
||
}
|
||
// Initialize ONNX Runtime session
|
||
instanceId_ = instanceCounter_.fetch_add(1); // Atomic increment
|
||
result = Init(_modelFilePath, true, 0);
|
||
_modelLoadValid = true;
|
||
_isInitialized = true;
|
||
return result;
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::LoadModel", e.what(), __FILE__, __LINE__);
|
||
return false;
|
||
}
|
||
}
|
||
bool ANSONNXOBB::LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap) {
|
||
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
||
ModelLoadingGuard mlg(_modelLoading);
|
||
try {
|
||
bool result = ANSODBase::LoadModelFromFolder(licenseKey, modelConfig, modelName, className, modelFolder, labelMap);
|
||
if (!result) return false;
|
||
std::string _modelName = modelName;
|
||
if (_modelName.empty()) {
|
||
_modelName = "train_last";
|
||
}
|
||
std::string modelFullName = _modelName + ".onnx";
|
||
// Parsing for YOLO only here
|
||
_modelConfig = modelConfig;
|
||
_modelConfig.detectionType = ANSCENTER::DetectionType::CLASSIFICATION;
|
||
_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
|
||
|
||
// 0. Check if the configuration file exist
|
||
if (FileExist(_modelConfigFile)) {
|
||
ModelType modelType;
|
||
std::vector<int> 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("ANSONNXOBB::Initialize. Load classes from string", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromString();
|
||
}
|
||
else {
|
||
this->_logger.LogDebug("ANSONNXOBB::Initialize. Load classes from file", _classFilePath, __FILE__, __LINE__);
|
||
LoadClassesFromFile();
|
||
}
|
||
}
|
||
// 1. Load labelMap and engine
|
||
labelMap.clear();
|
||
if (!_classes.empty())
|
||
labelMap = VectorToCommaSeparatedString(_classes);
|
||
// 2. Initialize ONNX Runtime session
|
||
instanceId_ = instanceCounter_.fetch_add(1); // Atomic increment
|
||
_modelLoadValid = true;
|
||
_isInitialized = true;
|
||
return result;
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::LoadModelFromFolder", e.what(), __FILE__, __LINE__);
|
||
return false;
|
||
}
|
||
}
|
||
std::vector<Object> ANSONNXOBB::RunInference(const cv::Mat& input, const std::string& camera_id) {
|
||
if (!PreInferenceCheck("ANSONNXOBB::RunInference")) return {};
|
||
try {
|
||
std::vector<Object> result;
|
||
if (input.empty()) return result;
|
||
if ((input.cols < 5) || (input.rows < 5)) return result;
|
||
result = detect(input, camera_id);
|
||
if (_trackerEnabled) {
|
||
result = ApplyTracking(result, camera_id);
|
||
if (_stabilizationEnabled) result = StabilizeDetections(result, camera_id);
|
||
}
|
||
return result;
|
||
}
|
||
catch (const std::exception& e) {
|
||
this->_logger.LogFatal("ANSONNXOBB::RunInference", e.what(), __FILE__, __LINE__);
|
||
return {};
|
||
}
|
||
}
|
||
std::vector<Object> ANSONNXOBB::RunInference(const cv::Mat& inputImgBGR) {
|
||
return RunInference(inputImgBGR, "CustomCam");
|
||
}
|
||
}
|