File size: 4,981 Bytes
e9eb33d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <algorithm>
#include <array>
#include <cmath>
#include <iostream>
#include <string>
#include <vector>
using namespace cv;
static std::string argVal(int argc, char** argv, const std::string& key, const std::string& def)
{
for (int i = 1; i + 1 < argc; ++i)
if (key == argv[i]) return argv[i + 1];
return def;
}
struct Det { float x1, y1, x2, y2, score; int cid; };
int main(int argc, char** argv)
{
std::string model = argVal(argc, argv, "--model", "efficientdet-d0_2026jul.onnx");
std::string image = argVal(argc, argv, "--image", "example_outputs/input_image.png");
std::string output = argVal(argc, argv, "--output", "example_outputs/output_image.png");
float conf = std::stof(argVal(argc, argv, "--conf", "0.4"));
const int sz = 512;
Mat img = imread(image);
if (img.empty()) { std::cerr << "could not read image: " << image << std::endl; return 1; }
Mat rgb;
cvtColor(img, rgb, COLOR_BGR2RGB);
resize(rgb, rgb, Size(sz, sz));
if (!rgb.isContinuous()) rgb = rgb.clone();
int blobShape[] = {1, sz, sz, 3};
Mat blob(4, blobShape, CV_8U, rgb.data);
dnn::Net net = dnn::readNetFromONNX(model);
net.setInput(blob);
std::vector<Mat> outs;
net.forward(outs, net.getUnconnectedOutLayersNames());
const float* boxp = nullptr;
const float* clsp = nullptr;
int n = 0, nc = 0;
for (size_t i = 0; i < outs.size(); ++i)
{
const Mat& o = outs[i];
const float* p = (const float*)o.data;
int last = o.size[o.dims - 1];
if (last == 4) { boxp = p; n = o.size[o.dims - 2]; }
else { clsp = p; nc = last; }
}
std::vector<std::array<float, 2>> baseWH;
double asp[3][2] = {{1.0, 1.0}, {1.4, 0.7}, {0.7, 1.4}};
for (int i = 0; i < 3; ++i) {
double s = std::pow(2.0, i / 3.0);
for (int a = 0; a < 3; ++a)
baseWH.push_back({(float)(32.0 * s * asp[a][0]), (float)(32.0 * s * asp[a][1])});
}
std::vector<float> acx, acy, aw, ah;
for (int lvl = 0; lvl < 5; ++lvl) {
int f = sz / (8 << lvl);
int step = 8 << lvl;
int m = 1 << lvl;
for (int y = 0; y < f; ++y)
for (int x = 0; x < f; ++x) {
float cx = (x + 0.5f) * step;
float cy = (y + 0.5f) * step;
for (auto& b : baseWH) {
acx.push_back(cx); acy.push_back(cy);
aw.push_back(b[0] * m); ah.push_back(b[1] * m);
}
}
}
std::vector<Det> dets;
for (int a = 0; a < n; ++a) {
const float* bp = boxp + (size_t)a * 4;
float ycenter = bp[0] * ah[a] + acy[a];
float xcenter = bp[1] * aw[a] + acx[a];
float bhv = std::exp(bp[2]) * ah[a];
float bwv = std::exp(bp[3]) * aw[a];
const float* cp = clsp + (size_t)a * nc;
int best = 0; float bestLogit = cp[0];
for (int c = 1; c < nc; ++c) if (cp[c] > bestLogit) { bestLogit = cp[c]; best = c; }
float score = 1.0f / (1.0f + std::exp(-bestLogit));
if (score > conf)
dets.push_back({(xcenter - bwv / 2) / sz, (ycenter - bhv / 2) / sz,
(xcenter + bwv / 2) / sz, (ycenter + bhv / 2) / sz, score, best});
}
std::sort(dets.begin(), dets.end(), [](const Det& a, const Det& b) { return a.score > b.score; });
std::vector<char> removed(dets.size(), 0);
std::vector<int> pick;
for (size_t i = 0; i < dets.size(); ++i) {
if (removed[i]) continue;
pick.push_back((int)i);
for (size_t j = i + 1; j < dets.size(); ++j) {
if (removed[j]) continue;
float xx1 = std::max(dets[i].x1, dets[j].x1);
float yy1 = std::max(dets[i].y1, dets[j].y1);
float xx2 = std::min(dets[i].x2, dets[j].x2);
float yy2 = std::min(dets[i].y2, dets[j].y2);
float inter = std::max(0.0f, xx2 - xx1) * std::max(0.0f, yy2 - yy1);
float ai = (dets[i].x2 - dets[i].x1) * (dets[i].y2 - dets[i].y1);
float aj = (dets[j].x2 - dets[j].x1) * (dets[j].y2 - dets[j].y1);
if (inter / (ai + aj - inter + 1e-9f) > 0.6f) removed[j] = 1;
}
}
std::cout << "efficientdet-d0 " << pick.size() << " detections" << std::endl;
int w = img.cols, h = img.rows;
for (int idx : pick) {
const Det& d = dets[idx];
std::cout << format("%d %.3f %.3f %.3f %.3f %.3f", d.cid, d.score, d.x1, d.y1, d.x2, d.y2) << std::endl;
rectangle(img, Point((int)(d.x1 * w), (int)(d.y1 * h)), Point((int)(d.x2 * w), (int)(d.y2 * h)), Scalar(0, 255, 0), 2);
putText(img, format("%d:%.2f", d.cid, d.score), Point((int)(d.x1 * w), (int)(d.y1 * h) - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 1);
}
imwrite(output, img);
std::cout << "wrote " << output << std::endl;
return 0;
}
|