|
#include "vp_yolo_detector_node.h"
|
|
namespace vp_nodes {
|
|
vp_yolo_detector_node::vp_yolo_detector_node(std::string node_name,
|
std::string model_path,
|
std::string model_config_path,
|
std::string labels_path,
|
int input_width,
|
int input_height,
|
int batch_size,
|
int class_id_offset,
|
float score_threshold,
|
float confidence_threshold,
|
float nms_threshold,
|
float scale,
|
cv::Scalar mean,
|
cv::Scalar std,
|
bool swap_rb):
|
vp_primary_infer_node(node_name, model_path, model_config_path,labels_path,input_width, input_height, batch_size, class_id_offset, scale, mean, std, swap_rb),
|
score_threshold(score_threshold),
|
confidence_threshold(confidence_threshold),
|
nms_threshold(nms_threshold) {
|
this->initialized();
|
}
|
|
vp_yolo_detector_node::~vp_yolo_detector_node() {
|
deinitialized();
|
}
|
|
void vp_yolo_detector_node::postprocess(const std::vector<cv::Mat>& raw_outputs, const std::vector<std::shared_ptr<vp_objects::vp_frame_meta>>& frame_meta_with_batch) {
|
// make sure heads of output are not zero
|
assert(raw_outputs.size() > 0);
|
|
// check dims of each output
|
auto dim_offset = raw_outputs[0].dims == 2 ? 0 : 1;
|
auto batch = dim_offset == 0 ? 1 : raw_outputs[0].size[0];
|
|
std::vector<int> class_ids;
|
std::vector<float> confidences;
|
std::vector<cv::Rect> boxes;
|
|
for (int k = 0; k < batch; k++) { // scan batch
|
auto& frame_meta = frame_meta_with_batch[k];
|
for (int i = 0; i < raw_outputs.size(); ++i) { // scan heads of output
|
cv::Mat output = cv::Mat(raw_outputs[i].size[0 + dim_offset], raw_outputs[i].size[1 + dim_offset], CV_32F, const_cast<uchar*>(raw_outputs[i].ptr(k)));
|
|
auto data = (float*)output.data;
|
for (int j = 0; j < output.rows; ++j, data += output.cols) {
|
|
float confidence = data[4];
|
// check confidence threshold
|
if (confidence < confidence_threshold) {
|
continue;
|
}
|
|
cv::Mat scores = output.row(j).colRange(5, output.cols);
|
cv::Point class_id;
|
double max_score;
|
// Get the value and location of the maximum score
|
cv::minMaxLoc(scores, 0, &max_score, 0, &class_id);
|
|
// check score threshold
|
if (max_score >= score_threshold) {
|
// smaller than 1.0 means relative value
|
// greater than 1.0 means absolute value which need to be converted to relative value
|
int centerX = (int)((data[0] <= 1 ? data[0] : data[0] / input_width) * frame_meta->frame.cols);
|
int centerY = (int)((data[1] <= 1 ? data[1] : data[1] / input_height) * frame_meta->frame.rows);
|
int width = (int)((data[2] <= 1 ? data[2] : data[2] / input_width) * frame_meta->frame.cols);
|
int height = (int)((data[3] <= 1 ? data[3] : data[3] / input_height) * frame_meta->frame.rows);
|
int left = centerX - width / 2;
|
int top = centerY - height / 2;
|
|
class_ids.push_back(class_id.x);
|
confidences.push_back((float)confidence);
|
boxes.push_back(cv::Rect(left, top, width, height));
|
}
|
}
|
}
|
|
// nms
|
std::vector<int> indices;
|
cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, nms_threshold, indices);
|
|
// create target
|
for (int i = 0; i < indices.size(); ++i) {
|
int idx = indices[i];
|
auto box = boxes[idx];
|
auto confidence = confidences[idx];
|
auto class_id = class_ids[idx];
|
auto label = (labels.size() < class_id + 1) ? "" : labels[class_id];
|
|
// check value range
|
box.x = std::max(box.x, 0);
|
box.y = std::max(box.y, 0);
|
box.width = std::min(box.width, frame_meta->frame.cols - box.x);
|
box.height = std::min(box.height, frame_meta->frame.rows - box.y);
|
if (box.width <= 0 || box.height <= 0) {
|
continue;
|
}
|
|
// apply offset to class id since multi detectors can exist at front of this one.
|
// later we MUST use the new class id (applied offset) instead of orignal one.
|
class_id += class_id_offset;
|
|
auto target = std::make_shared<vp_objects::vp_frame_target>(box.x, box.y, box.width, box.height, class_id, confidence, frame_meta->frame_index, frame_meta->channel_index, label);
|
|
// insert target back to frame meta
|
frame_meta->targets.push_back(target);
|
}
|
|
class_ids.clear();
|
confidences.clear();
|
boxes.clear();
|
}
|
}
|
}
|
