|
| 1 | +#pragma once |
| 2 | + |
| 3 | +#include "YoloONNX.hpp" |
| 4 | + |
| 5 | +/// |
| 6 | +/// \brief The DFINE_is_onnx class |
| 7 | +/// |
| 8 | +class DFINE_is_onnx : public YoloONNX |
| 9 | +{ |
| 10 | +public: |
| 11 | + DFINE_is_onnx(std::vector<std::string>& inputTensorNames, std::vector<std::string>& outputTensorNames) |
| 12 | + { |
| 13 | + inputTensorNames.push_back("input"); |
| 14 | + outputTensorNames.push_back("logits"); |
| 15 | + outputTensorNames.push_back("boxes"); |
| 16 | + outputTensorNames.push_back("mask_probs"); |
| 17 | + } |
| 18 | + |
| 19 | +protected: |
| 20 | + /// |
| 21 | + /// \brief GetResult |
| 22 | + /// \param output |
| 23 | + /// \return |
| 24 | + /// |
| 25 | + std::vector<tensor_rt::Result> GetResult(size_t imgIdx, int /*keep_topk*/, const std::vector<float*>& outputs, cv::Size frameSize) |
| 26 | + { |
| 27 | + std::vector<tensor_rt::Result> resBoxes; |
| 28 | + |
| 29 | + //0: name: input, size: 1x3x640x640 |
| 30 | + //1: name: logits, size: 1x300x80 |
| 31 | + //2: name: boxes, size: 1x300x4 |
| 32 | + //3: name: mask_probs, size: 1x300x160x160 |
| 33 | + |
| 34 | + |
| 35 | + //0: name: input, size: 1x3x432x432 |
| 36 | + //1: name: dets, size: 1x200x4 |
| 37 | + //2: name: labels, size: 1x200x91 |
| 38 | + //3: name: 4245, size: 1x200x108x108 |
| 39 | + |
| 40 | + const float fw = static_cast<float>(frameSize.width) / static_cast<float>(m_resizedROI.width); |
| 41 | + const float fh = static_cast<float>(frameSize.height) / static_cast<float>(m_resizedROI.height); |
| 42 | + |
| 43 | + cv::Size inputSize(m_inputDims[0].d[3], m_inputDims[0].d[2]); |
| 44 | + cv::Size2f inputSizef(static_cast<float>(inputSize.width), static_cast<float>(inputSize.height)); |
| 45 | + |
| 46 | + //std::cout << "m_resizedROI: " << m_resizedROI << ", frameSize: " << frameSize << ", fw_h: " << cv::Size2f(fw, fh) << ", m_inputDims: " << cv::Point3i(m_inputDims.d[1], m_inputDims.d[2], m_inputDims.d[3]) << std::endl; |
| 47 | + |
| 48 | + int labelsInd = 0; |
| 49 | + int detsInd = 1; |
| 50 | + int segInd = 2; |
| 51 | + |
| 52 | + auto dets = outputs[detsInd]; |
| 53 | + auto labels = outputs[labelsInd]; |
| 54 | + |
| 55 | + auto masks = outputs[segInd]; |
| 56 | + |
| 57 | + size_t ncInd = 2; |
| 58 | + size_t lenInd = 1; |
| 59 | + |
| 60 | + |
| 61 | + size_t nc = m_outpuDims[labelsInd].d[ncInd]; |
| 62 | + size_t len = static_cast<size_t>(m_outpuDims[detsInd].d[lenInd]) / m_params.m_explicitBatchSize; |
| 63 | + auto volume0 = len * m_outpuDims[detsInd].d[ncInd]; // Volume(m_outpuDims[0]); |
| 64 | + dets += volume0 * imgIdx; |
| 65 | + auto volume1 = len * m_outpuDims[labelsInd].d[ncInd]; // Volume(m_outpuDims[0]); |
| 66 | + labels += volume1 * imgIdx; |
| 67 | + |
| 68 | + int segChannels = static_cast<int>(m_outpuDims[segInd].d[1]); |
| 69 | + int segWidth = static_cast<int>(m_outpuDims[segInd].d[2]); |
| 70 | + int segHeight = static_cast<int>(m_outpuDims[segInd].d[3]); |
| 71 | + masks += imgIdx * segChannels * segWidth * segHeight; |
| 72 | + |
| 73 | + cv::Mat binaryMask8U(segHeight, segWidth, CV_8UC1); |
| 74 | + |
| 75 | + //std::cout << "len = " << len << ", nc = " << nc << ", m_params.confThreshold = " << m_params.m_confThreshold << ", volume0 = " << volume0 << ", volume1 = " << volume1 << std::endl; |
| 76 | + |
| 77 | + auto L2Conf = [](float v) |
| 78 | + { |
| 79 | + return 1.f / (1.f + std::exp(-v)); |
| 80 | + }; |
| 81 | + |
| 82 | + for (size_t i = 0; i < len; ++i) |
| 83 | + { |
| 84 | + float classConf = L2Conf(labels[0]); |
| 85 | + size_t classId = 0; |
| 86 | + for (size_t cli = 1; cli < nc; ++cli) |
| 87 | + { |
| 88 | + auto conf = L2Conf(labels[cli]); |
| 89 | + if (classConf < conf) |
| 90 | + { |
| 91 | + classConf = conf; |
| 92 | + classId = cli; |
| 93 | + } |
| 94 | + } |
| 95 | + |
| 96 | + if (classConf >= m_params.m_confThreshold) |
| 97 | + { |
| 98 | + float d0 = dets[0]; |
| 99 | + float d1 = dets[1]; |
| 100 | + float d2 = dets[2]; |
| 101 | + float d3 = dets[3]; |
| 102 | + |
| 103 | + float x = fw * (inputSizef.width * (d0 - d2 / 2.f) - m_resizedROI.x); |
| 104 | + float y = fh * (inputSizef.height * (d1 - d3 / 2.f) - m_resizedROI.y); |
| 105 | + float width = fw * inputSizef.width * d2; |
| 106 | + float height = fh * inputSizef.height * d3; |
| 107 | + |
| 108 | + //if (i == 0) |
| 109 | + //{ |
| 110 | + // std::cout << i << ": classConf = " << classConf << ", classId = " << classId << " (" << labels[classId] << "), rect = " << cv::Rect2f(x, y, width, height) << std::endl; |
| 111 | + // std::cout << "dets = " << d0 << ", " << d1 << ", " << d2 << ", " << d3 << std::endl; |
| 112 | + //} |
| 113 | + resBoxes.emplace_back(classId, classConf, cv::Rect(cvRound(x), cvRound(y), cvRound(width), cvRound(height))); |
| 114 | + |
| 115 | + double maskThreshold = 0.1; |
| 116 | + for (int row = 0; row < segHeight; ++row) |
| 117 | + { |
| 118 | + const float* maskPtr = masks + row * segWidth; |
| 119 | + uchar* binMaskPtr = binaryMask8U.ptr(row); |
| 120 | + |
| 121 | + for (int col = 0; col < segWidth; ++col) |
| 122 | + { |
| 123 | + binMaskPtr[col] = (maskPtr[col] > maskThreshold) ? 255 : 0; |
| 124 | + } |
| 125 | + } |
| 126 | + |
| 127 | + tensor_rt::Result& resObj = resBoxes.back(); |
| 128 | + |
| 129 | + cv::Rect smallRect; |
| 130 | + smallRect.x = cvRound(segHeight * (d0 - d2 / 2.f)); |
| 131 | + smallRect.y = cvRound(segHeight * (d1 - d3 / 2.f)); |
| 132 | + smallRect.width = cvRound(segHeight * d2); |
| 133 | + smallRect.height = cvRound(segHeight * d3); |
| 134 | + smallRect = Clamp(smallRect, cv::Size(segWidth, segHeight)); |
| 135 | + |
| 136 | + if (smallRect.area() > 0) |
| 137 | + { |
| 138 | + cv::resize(binaryMask8U(smallRect), resObj.m_boxMask, resObj.m_brect.size(), 0, 0, cv::INTER_NEAREST); |
| 139 | + |
| 140 | +#if 0 |
| 141 | + static int globalObjInd = 0; |
| 142 | + SaveMat(mask, std::to_string(globalObjInd) + "_mask", ".png", "tmp", true); |
| 143 | + SaveMat(binaryMask, std::to_string(globalObjInd) + "_bin_mask", ".png", "tmp", true); |
| 144 | + SaveMat(binaryMask8U, std::to_string(globalObjInd) + "_bin_mask_8u", ".png", "tmp", true); |
| 145 | + SaveMat(resObj.m_boxMask, std::to_string(globalObjInd++) + "_obj_mask", ".png", "tmp", true); |
| 146 | + std::cout << "inputSize: " << inputSize << ", localRect: " << localRect << std::endl; |
| 147 | +#endif |
| 148 | + |
| 149 | + std::vector<std::vector<cv::Point>> contours; |
| 150 | + std::vector<cv::Vec4i> hierarchy; |
| 151 | + cv::findContours(resObj.m_boxMask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE, cv::Point()); |
| 152 | + |
| 153 | + for (const auto& contour : contours) |
| 154 | + { |
| 155 | + cv::Rect br = cv::boundingRect(contour); |
| 156 | + |
| 157 | + if (br.width >= 4 && |
| 158 | + br.height >= 4) |
| 159 | + { |
| 160 | + int dx = resObj.m_brect.x; |
| 161 | + int dy = resObj.m_brect.y; |
| 162 | + |
| 163 | + cv::RotatedRect rr = (contour.size() < 5) ? cv::minAreaRect(contour) : cv::fitEllipse(contour); |
| 164 | + rr.center.x = rr.center.x * fw + dx; |
| 165 | + rr.center.y = rr.center.y * fw + dy; |
| 166 | + rr.size.width *= fw; |
| 167 | + rr.size.height *= fh; |
| 168 | + |
| 169 | + br.x = cvRound(dx + br.x * fw); |
| 170 | + br.y = cvRound(dy + br.y * fh); |
| 171 | + br.width = cvRound(br.width * fw); |
| 172 | + br.height = cvRound(br.height * fh); |
| 173 | + |
| 174 | + resObj.m_brect = br; |
| 175 | + //resObj.m_rrect = rr; |
| 176 | + |
| 177 | + //std::cout << "resBoxes[" << i << "] br: " << br << ", rr: (" << rr.size << " from " << rr.center << ", " << rr.angle << ")" << std::endl; |
| 178 | + |
| 179 | + break; |
| 180 | + } |
| 181 | + } |
| 182 | + } |
| 183 | + else |
| 184 | + { |
| 185 | + resObj.m_boxMask = cv::Mat(resObj.m_brect.size(), CV_8UC1, cv::Scalar(255)); |
| 186 | + } |
| 187 | + } |
| 188 | + |
| 189 | + dets += m_outpuDims[detsInd].d[ncInd]; |
| 190 | + labels += m_outpuDims[labelsInd].d[ncInd]; |
| 191 | + masks += segWidth * segHeight; |
| 192 | + } |
| 193 | + |
| 194 | + return resBoxes; |
| 195 | + } |
| 196 | +}; |
0 commit comments