|
| 1 | +#include "YoloDetector.h" |
| 2 | +#include "nms.h" |
| 3 | + |
| 4 | +/// |
| 5 | +/// \brief YoloDetector::YoloDetector |
| 6 | +/// \param collectPoints |
| 7 | +/// \param gray |
| 8 | +/// |
| 9 | +YoloDetector::YoloDetector( |
| 10 | + bool collectPoints, |
| 11 | + cv::UMat& colorFrame |
| 12 | + ) |
| 13 | + : |
| 14 | + BaseDetector(collectPoints, colorFrame), |
| 15 | + m_WHRatio(InWidth / (float)InHeight), |
| 16 | + m_inScaleFactor(0.003921f), |
| 17 | + m_meanVal(0), |
| 18 | + m_confidenceThreshold(0.24f), |
| 19 | + m_maxCropRatio(2.0f) |
| 20 | +{ |
| 21 | + m_classNames = { "background", |
| 22 | + "aeroplane", "bicycle", "bird", "boat", |
| 23 | + "bottle", "bus", "car", "cat", "chair", |
| 24 | + "cow", "diningtable", "dog", "horse", |
| 25 | + "motorbike", "person", "pottedplant", |
| 26 | + "sheep", "sofa", "train", "tvmonitor" }; |
| 27 | +} |
| 28 | + |
| 29 | +/// |
| 30 | +/// \brief YoloDetector::~YoloDetector |
| 31 | +/// |
| 32 | +YoloDetector::~YoloDetector(void) |
| 33 | +{ |
| 34 | +} |
| 35 | + |
| 36 | +/// |
| 37 | +/// \brief YoloDetector::Init |
| 38 | +/// \return |
| 39 | +/// |
| 40 | +bool YoloDetector::Init(const config_t& config) |
| 41 | +{ |
| 42 | + auto modelConfiguration = config.find("modelConfiguration"); |
| 43 | + auto modelBinary = config.find("modelBinary"); |
| 44 | + if (modelConfiguration != config.end() && modelBinary != config.end()) |
| 45 | + { |
| 46 | + m_net = cv::dnn::readNetFromDarknet(modelConfiguration->second, modelBinary->second); |
| 47 | + } |
| 48 | + |
| 49 | + auto classNames = config.find("classNames"); |
| 50 | + if (classNames != config.end()) |
| 51 | + { |
| 52 | + std::ifstream classNamesFile(classNames->second); |
| 53 | + if (classNamesFile.is_open()) |
| 54 | + { |
| 55 | + m_classNames.clear(); |
| 56 | + std::string className; |
| 57 | + for (; std::getline(classNamesFile, className); ) |
| 58 | + { |
| 59 | + m_classNames.push_back(className); |
| 60 | + } |
| 61 | + } |
| 62 | + } |
| 63 | + |
| 64 | + auto confidenceThreshold = config.find("confidenceThreshold"); |
| 65 | + if (confidenceThreshold != config.end()) |
| 66 | + { |
| 67 | + m_confidenceThreshold = std::stof(confidenceThreshold->second); |
| 68 | + } |
| 69 | + |
| 70 | + auto maxCropRatio = config.find("maxCropRatio"); |
| 71 | + if (maxCropRatio != config.end()) |
| 72 | + { |
| 73 | + m_maxCropRatio = std::stof(maxCropRatio->second); |
| 74 | + if (m_maxCropRatio < 1.f) |
| 75 | + { |
| 76 | + m_maxCropRatio = 1.f; |
| 77 | + } |
| 78 | + } |
| 79 | + |
| 80 | + return !m_net.empty(); |
| 81 | +} |
| 82 | + |
| 83 | +/// |
| 84 | +/// \brief YoloDetector::Detect |
| 85 | +/// \param gray |
| 86 | +/// |
| 87 | +void YoloDetector::Detect(cv::UMat& colorFrame) |
| 88 | +{ |
| 89 | + m_regions.clear(); |
| 90 | + |
| 91 | + regions_t tmpRegions; |
| 92 | + |
| 93 | + cv::Mat colorMat = colorFrame.getMat(cv::ACCESS_READ); |
| 94 | + |
| 95 | + int cropHeight = cvRound(m_maxCropRatio * InHeight); |
| 96 | + int cropWidth = cvRound(m_maxCropRatio * InWidth); |
| 97 | + |
| 98 | + if (colorFrame.cols / (float)colorFrame.rows > m_WHRatio) |
| 99 | + { |
| 100 | + if (m_maxCropRatio <= 0 || cropHeight >= colorFrame.rows) |
| 101 | + { |
| 102 | + cropHeight = colorFrame.rows; |
| 103 | + } |
| 104 | + cropWidth = cvRound(cropHeight * m_WHRatio); |
| 105 | + } |
| 106 | + else |
| 107 | + { |
| 108 | + if (m_maxCropRatio <= 0 || cropWidth >= colorFrame.cols) |
| 109 | + { |
| 110 | + cropWidth = colorFrame.cols; |
| 111 | + } |
| 112 | + cropHeight = cvRound(colorFrame.cols / m_WHRatio); |
| 113 | + } |
| 114 | + |
| 115 | + cv::Rect crop(0, 0, cropWidth, cropHeight); |
| 116 | + |
| 117 | + for (; crop.y < colorMat.rows; crop.y += crop.height / 2) |
| 118 | + { |
| 119 | + bool needBreakY = false; |
| 120 | + if (crop.y + crop.height >= colorMat.rows) |
| 121 | + { |
| 122 | + crop.y = colorMat.rows - crop.height; |
| 123 | + needBreakY = true; |
| 124 | + } |
| 125 | + for (crop.x = 0; crop.x < colorMat.cols; crop.x += crop.width / 2) |
| 126 | + { |
| 127 | + bool needBreakX = false; |
| 128 | + if (crop.x + crop.width >= colorMat.cols) |
| 129 | + { |
| 130 | + crop.x = colorMat.cols - crop.width; |
| 131 | + needBreakX = true; |
| 132 | + } |
| 133 | + |
| 134 | + DetectInCrop(colorMat, crop, tmpRegions); |
| 135 | + |
| 136 | + if (needBreakX) |
| 137 | + { |
| 138 | + break; |
| 139 | + } |
| 140 | + } |
| 141 | + if (needBreakY) |
| 142 | + { |
| 143 | + break; |
| 144 | + } |
| 145 | + } |
| 146 | + |
| 147 | + nms3<CRegion>(tmpRegions, m_regions, 0.4f, |
| 148 | + [](const CRegion& reg) -> cv::Rect { return reg.m_rect; }, |
| 149 | + [](const CRegion& reg) -> float { return reg.m_confidence; }, |
| 150 | + 0, 0.f); |
| 151 | + |
| 152 | + if (m_collectPoints) |
| 153 | + { |
| 154 | + for (auto& region : m_regions) |
| 155 | + { |
| 156 | + CollectPoints(region); |
| 157 | + } |
| 158 | + } |
| 159 | +} |
| 160 | + |
| 161 | +/// |
| 162 | +/// \brief YoloDetector::DetectInCrop |
| 163 | +/// \param colorFrame |
| 164 | +/// \param crop |
| 165 | +/// \param tmpRegions |
| 166 | +/// |
| 167 | +void YoloDetector::DetectInCrop(cv::Mat colorFrame, const cv::Rect& crop, regions_t& tmpRegions) |
| 168 | +{ |
| 169 | + //Convert Mat to batch of images |
| 170 | + cv::Mat inputBlob = cv::dnn::blobFromImage(cv::Mat(colorFrame, crop), m_inScaleFactor, cv::Size(InWidth, InHeight), m_meanVal, false, true); |
| 171 | + |
| 172 | + m_net.setInput(inputBlob, "data"); //set the network input |
| 173 | + |
| 174 | + cv::Mat detectionMat = m_net.forward("detection_out"); //compute output |
| 175 | + |
| 176 | + for (int i = 0; i < detectionMat.rows; ++i) |
| 177 | + { |
| 178 | + const int probability_index = 5; |
| 179 | + const int probability_size = detectionMat.cols - probability_index; |
| 180 | + float* prob_array_ptr = &detectionMat.at<float>(i, probability_index); |
| 181 | + |
| 182 | + size_t objectClass = std::max_element(prob_array_ptr, prob_array_ptr + probability_size) - prob_array_ptr; |
| 183 | + float confidence = detectionMat.at<float>(i, (int)objectClass + probability_index); |
| 184 | + |
| 185 | + if (confidence > m_confidenceThreshold) |
| 186 | + { |
| 187 | + float x_center = detectionMat.at<float>(i, 0) * crop.width + crop.x; |
| 188 | + float y_center = detectionMat.at<float>(i, 1) * crop.height + crop.y; |
| 189 | + float width = detectionMat.at<float>(i, 2) * crop.width; |
| 190 | + float height = detectionMat.at<float>(i, 3) * crop.height; |
| 191 | + cv::Point p1(cvRound(x_center - width / 2), cvRound(y_center - height / 2)); |
| 192 | + cv::Point p2(cvRound(x_center + width / 2), cvRound(y_center + height / 2)); |
| 193 | + cv::Rect object(p1, p2); |
| 194 | + |
| 195 | + tmpRegions.push_back(CRegion(object, m_classNames[objectClass], confidence)); |
| 196 | + } |
| 197 | + } |
| 198 | +} |
0 commit comments