398 lines
15 KiB
C++
398 lines
15 KiB
C++
#include <atomic>
|
|
#include <chrono>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
#include "node.h"
|
|
|
|
namespace rk3588 {
|
|
|
|
namespace {
|
|
|
|
constexpr int kObjClassNum = 80;
|
|
|
|
const char* kCocoLabels[kObjClassNum] = {
|
|
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
|
|
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
|
|
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
|
|
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
|
|
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
|
|
"bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
|
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
|
|
"couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
|
|
"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
|
|
"book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
|
|
};
|
|
|
|
struct Color {
|
|
uint8_t r, g, b;
|
|
};
|
|
|
|
const Color kClassColors[] = {
|
|
{255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255},
|
|
{0, 255, 255}, {128, 0, 0}, {0, 128, 0}, {0, 0, 128}, {128, 128, 0},
|
|
{128, 0, 128}, {0, 128, 128}, {255, 128, 0}, {255, 0, 128}, {128, 255, 0},
|
|
{0, 255, 128}, {128, 0, 255}, {0, 128, 255}, {255, 128, 128}, {128, 255, 128}
|
|
};
|
|
|
|
inline Color GetClassColor(int cls_id) {
|
|
return kClassColors[cls_id % 20];
|
|
}
|
|
|
|
inline const char* GetClassName(int cls_id) {
|
|
if (cls_id >= 0 && cls_id < kObjClassNum) {
|
|
return kCocoLabels[cls_id];
|
|
}
|
|
return "unknown";
|
|
}
|
|
|
|
inline int Clamp(int val, int min_val, int max_val) {
|
|
return val < min_val ? min_val : (val > max_val ? max_val : val);
|
|
}
|
|
|
|
void DrawHLine(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
|
int x1, int x2, int y, int thickness, const Color& color) {
|
|
x1 = Clamp(x1, 0, w - 1);
|
|
x2 = Clamp(x2, 0, w - 1);
|
|
if (x1 > x2) std::swap(x1, x2);
|
|
|
|
for (int t = 0; t < thickness; ++t) {
|
|
int cy = y + t;
|
|
if (cy < 0 || cy >= h) continue;
|
|
|
|
if (fmt == PixelFormat::RGB) {
|
|
for (int x = x1; x <= x2; ++x) {
|
|
int idx = (cy * stride) + x * 3;
|
|
data[idx] = color.r;
|
|
data[idx + 1] = color.g;
|
|
data[idx + 2] = color.b;
|
|
}
|
|
} else if (fmt == PixelFormat::BGR) {
|
|
for (int x = x1; x <= x2; ++x) {
|
|
int idx = (cy * stride) + x * 3;
|
|
data[idx] = color.b;
|
|
data[idx + 1] = color.g;
|
|
data[idx + 2] = color.r;
|
|
}
|
|
} else if (fmt == PixelFormat::NV12) {
|
|
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
|
for (int x = x1; x <= x2; ++x) {
|
|
data[cy * w + x] = Y;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void DrawVLine(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
|
int x, int y1, int y2, int thickness, const Color& color) {
|
|
y1 = Clamp(y1, 0, h - 1);
|
|
y2 = Clamp(y2, 0, h - 1);
|
|
if (y1 > y2) std::swap(y1, y2);
|
|
|
|
for (int t = 0; t < thickness; ++t) {
|
|
int cx = x + t;
|
|
if (cx < 0 || cx >= w) continue;
|
|
|
|
if (fmt == PixelFormat::RGB) {
|
|
for (int y = y1; y <= y2; ++y) {
|
|
int idx = (y * stride) + cx * 3;
|
|
data[idx] = color.r;
|
|
data[idx + 1] = color.g;
|
|
data[idx + 2] = color.b;
|
|
}
|
|
} else if (fmt == PixelFormat::BGR) {
|
|
for (int y = y1; y <= y2; ++y) {
|
|
int idx = (y * stride) + cx * 3;
|
|
data[idx] = color.b;
|
|
data[idx + 1] = color.g;
|
|
data[idx + 2] = color.r;
|
|
}
|
|
} else if (fmt == PixelFormat::NV12) {
|
|
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
|
for (int y = y1; y <= y2; ++y) {
|
|
data[y * w + cx] = Y;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void DrawRect(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
|
int x1, int y1, int x2, int y2, int thickness, const Color& color) {
|
|
DrawHLine(data, w, h, stride, fmt, x1, x2, y1, thickness, color);
|
|
DrawHLine(data, w, h, stride, fmt, x1, x2, y2 - thickness + 1, thickness, color);
|
|
DrawVLine(data, w, h, stride, fmt, x1, y1, y2, thickness, color);
|
|
DrawVLine(data, w, h, stride, fmt, x2 - thickness + 1, y1, y2, thickness, color);
|
|
}
|
|
|
|
const uint8_t kFont5x7[96][7] = {
|
|
{0x00,0x00,0x00,0x00,0x00,0x00,0x00}, // ' '
|
|
{0x04,0x04,0x04,0x04,0x00,0x00,0x04}, // '!'
|
|
{0x0A,0x0A,0x0A,0x00,0x00,0x00,0x00}, // '"'
|
|
{0x0A,0x0A,0x1F,0x0A,0x1F,0x0A,0x0A}, // '#'
|
|
{0x04,0x0F,0x14,0x0E,0x05,0x1E,0x04}, // '$'
|
|
{0x18,0x19,0x02,0x04,0x08,0x13,0x03}, // '%'
|
|
{0x0C,0x12,0x14,0x08,0x15,0x12,0x0D}, // '&'
|
|
{0x0C,0x04,0x08,0x00,0x00,0x00,0x00}, // '''
|
|
{0x02,0x04,0x08,0x08,0x08,0x04,0x02}, // '('
|
|
{0x08,0x04,0x02,0x02,0x02,0x04,0x08}, // ')'
|
|
{0x00,0x04,0x15,0x0E,0x15,0x04,0x00}, // '*'
|
|
{0x00,0x04,0x04,0x1F,0x04,0x04,0x00}, // '+'
|
|
{0x00,0x00,0x00,0x00,0x0C,0x04,0x08}, // ','
|
|
{0x00,0x00,0x00,0x1F,0x00,0x00,0x00}, // '-'
|
|
{0x00,0x00,0x00,0x00,0x00,0x0C,0x0C}, // '.'
|
|
{0x00,0x01,0x02,0x04,0x08,0x10,0x00}, // '/'
|
|
{0x0E,0x11,0x13,0x15,0x19,0x11,0x0E}, // '0'
|
|
{0x04,0x0C,0x04,0x04,0x04,0x04,0x0E}, // '1'
|
|
{0x0E,0x11,0x01,0x02,0x04,0x08,0x1F}, // '2'
|
|
{0x1F,0x02,0x04,0x02,0x01,0x11,0x0E}, // '3'
|
|
{0x02,0x06,0x0A,0x12,0x1F,0x02,0x02}, // '4'
|
|
{0x1F,0x10,0x1E,0x01,0x01,0x11,0x0E}, // '5'
|
|
{0x06,0x08,0x10,0x1E,0x11,0x11,0x0E}, // '6'
|
|
{0x1F,0x01,0x02,0x04,0x08,0x08,0x08}, // '7'
|
|
{0x0E,0x11,0x11,0x0E,0x11,0x11,0x0E}, // '8'
|
|
{0x0E,0x11,0x11,0x0F,0x01,0x02,0x0C}, // '9'
|
|
{0x00,0x0C,0x0C,0x00,0x0C,0x0C,0x00}, // ':'
|
|
{0x00,0x0C,0x0C,0x00,0x0C,0x04,0x08}, // ';'
|
|
{0x02,0x04,0x08,0x10,0x08,0x04,0x02}, // '<'
|
|
{0x00,0x00,0x1F,0x00,0x1F,0x00,0x00}, // '='
|
|
{0x08,0x04,0x02,0x01,0x02,0x04,0x08}, // '>'
|
|
{0x0E,0x11,0x01,0x02,0x04,0x00,0x04}, // '?'
|
|
{0x0E,0x11,0x17,0x15,0x17,0x10,0x0E}, // '@'
|
|
{0x0E,0x11,0x11,0x1F,0x11,0x11,0x11}, // 'A'
|
|
{0x1E,0x11,0x11,0x1E,0x11,0x11,0x1E}, // 'B'
|
|
{0x0E,0x11,0x10,0x10,0x10,0x11,0x0E}, // 'C'
|
|
{0x1C,0x12,0x11,0x11,0x11,0x12,0x1C}, // 'D'
|
|
{0x1F,0x10,0x10,0x1E,0x10,0x10,0x1F}, // 'E'
|
|
{0x1F,0x10,0x10,0x1E,0x10,0x10,0x10}, // 'F'
|
|
{0x0E,0x11,0x10,0x17,0x11,0x11,0x0F}, // 'G'
|
|
{0x11,0x11,0x11,0x1F,0x11,0x11,0x11}, // 'H'
|
|
{0x0E,0x04,0x04,0x04,0x04,0x04,0x0E}, // 'I'
|
|
{0x07,0x02,0x02,0x02,0x02,0x12,0x0C}, // 'J'
|
|
{0x11,0x12,0x14,0x18,0x14,0x12,0x11}, // 'K'
|
|
{0x10,0x10,0x10,0x10,0x10,0x10,0x1F}, // 'L'
|
|
{0x11,0x1B,0x15,0x15,0x11,0x11,0x11}, // 'M'
|
|
{0x11,0x11,0x19,0x15,0x13,0x11,0x11}, // 'N'
|
|
{0x0E,0x11,0x11,0x11,0x11,0x11,0x0E}, // 'O'
|
|
{0x1E,0x11,0x11,0x1E,0x10,0x10,0x10}, // 'P'
|
|
{0x0E,0x11,0x11,0x11,0x15,0x12,0x0D}, // 'Q'
|
|
{0x1E,0x11,0x11,0x1E,0x14,0x12,0x11}, // 'R'
|
|
{0x0F,0x10,0x10,0x0E,0x01,0x01,0x1E}, // 'S'
|
|
{0x1F,0x04,0x04,0x04,0x04,0x04,0x04}, // 'T'
|
|
{0x11,0x11,0x11,0x11,0x11,0x11,0x0E}, // 'U'
|
|
{0x11,0x11,0x11,0x11,0x11,0x0A,0x04}, // 'V'
|
|
{0x11,0x11,0x11,0x15,0x15,0x15,0x0A}, // 'W'
|
|
{0x11,0x11,0x0A,0x04,0x0A,0x11,0x11}, // 'X'
|
|
{0x11,0x11,0x11,0x0A,0x04,0x04,0x04}, // 'Y'
|
|
{0x1F,0x01,0x02,0x04,0x08,0x10,0x1F}, // 'Z'
|
|
{0x0E,0x08,0x08,0x08,0x08,0x08,0x0E}, // '['
|
|
{0x00,0x10,0x08,0x04,0x02,0x01,0x00}, // '\'
|
|
{0x0E,0x02,0x02,0x02,0x02,0x02,0x0E}, // ']'
|
|
{0x04,0x0A,0x11,0x00,0x00,0x00,0x00}, // '^'
|
|
{0x00,0x00,0x00,0x00,0x00,0x00,0x1F}, // '_'
|
|
{0x08,0x04,0x02,0x00,0x00,0x00,0x00}, // '`'
|
|
{0x00,0x00,0x0E,0x01,0x0F,0x11,0x0F}, // 'a'
|
|
{0x10,0x10,0x16,0x19,0x11,0x11,0x1E}, // 'b'
|
|
{0x00,0x00,0x0E,0x10,0x10,0x11,0x0E}, // 'c'
|
|
{0x01,0x01,0x0D,0x13,0x11,0x11,0x0F}, // 'd'
|
|
{0x00,0x00,0x0E,0x11,0x1F,0x10,0x0E}, // 'e'
|
|
{0x06,0x09,0x08,0x1C,0x08,0x08,0x08}, // 'f'
|
|
{0x00,0x0F,0x11,0x11,0x0F,0x01,0x0E}, // 'g'
|
|
{0x10,0x10,0x16,0x19,0x11,0x11,0x11}, // 'h'
|
|
{0x04,0x00,0x0C,0x04,0x04,0x04,0x0E}, // 'i'
|
|
{0x02,0x00,0x06,0x02,0x02,0x12,0x0C}, // 'j'
|
|
{0x10,0x10,0x12,0x14,0x18,0x14,0x12}, // 'k'
|
|
{0x0C,0x04,0x04,0x04,0x04,0x04,0x0E}, // 'l'
|
|
{0x00,0x00,0x1A,0x15,0x15,0x11,0x11}, // 'm'
|
|
{0x00,0x00,0x16,0x19,0x11,0x11,0x11}, // 'n'
|
|
{0x00,0x00,0x0E,0x11,0x11,0x11,0x0E}, // 'o'
|
|
{0x00,0x00,0x1E,0x11,0x1E,0x10,0x10}, // 'p'
|
|
{0x00,0x00,0x0D,0x13,0x0F,0x01,0x01}, // 'q'
|
|
{0x00,0x00,0x16,0x19,0x10,0x10,0x10}, // 'r'
|
|
{0x00,0x00,0x0E,0x10,0x0E,0x01,0x1E}, // 's'
|
|
{0x08,0x08,0x1C,0x08,0x08,0x09,0x06}, // 't'
|
|
{0x00,0x00,0x11,0x11,0x11,0x13,0x0D}, // 'u'
|
|
{0x00,0x00,0x11,0x11,0x11,0x0A,0x04}, // 'v'
|
|
{0x00,0x00,0x11,0x11,0x15,0x15,0x0A}, // 'w'
|
|
{0x00,0x00,0x11,0x0A,0x04,0x0A,0x11}, // 'x'
|
|
{0x00,0x00,0x11,0x11,0x0F,0x01,0x0E}, // 'y'
|
|
{0x00,0x00,0x1F,0x02,0x04,0x08,0x1F}, // 'z'
|
|
{0x02,0x04,0x04,0x08,0x04,0x04,0x02}, // '{'
|
|
{0x04,0x04,0x04,0x04,0x04,0x04,0x04}, // '|'
|
|
{0x08,0x04,0x04,0x02,0x04,0x04,0x08}, // '}'
|
|
{0x00,0x00,0x08,0x15,0x02,0x00,0x00}, // '~'
|
|
{0x00,0x00,0x00,0x00,0x00,0x00,0x00}, // DEL
|
|
};
|
|
|
|
void DrawChar(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
|
int x, int y, char c, int scale, const Color& color) {
|
|
if (c < 32 || c > 127) c = ' ';
|
|
int idx = c - 32;
|
|
const uint8_t* glyph = kFont5x7[idx];
|
|
|
|
for (int row = 0; row < 7; ++row) {
|
|
for (int col = 0; col < 5; ++col) {
|
|
if (glyph[row] & (1 << (4 - col))) {
|
|
for (int sy = 0; sy < scale; ++sy) {
|
|
for (int sx = 0; sx < scale; ++sx) {
|
|
int px = x + col * scale + sx;
|
|
int py = y + row * scale + sy;
|
|
if (px < 0 || px >= w || py < 0 || py >= h) continue;
|
|
|
|
if (fmt == PixelFormat::RGB) {
|
|
int i = py * stride + px * 3;
|
|
data[i] = color.r;
|
|
data[i + 1] = color.g;
|
|
data[i + 2] = color.b;
|
|
} else if (fmt == PixelFormat::BGR) {
|
|
int i = py * stride + px * 3;
|
|
data[i] = color.b;
|
|
data[i + 1] = color.g;
|
|
data[i + 2] = color.r;
|
|
} else if (fmt == PixelFormat::NV12) {
|
|
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
|
data[py * w + px] = Y;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void DrawText(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
|
int x, int y, const char* text, int scale, const Color& color) {
|
|
int cx = x;
|
|
while (*text) {
|
|
DrawChar(data, w, h, stride, fmt, cx, y, *text, scale, color);
|
|
cx += 6 * scale;
|
|
++text;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
class OsdNode : public INode {
|
|
public:
|
|
std::string Id() const override { return id_; }
|
|
std::string Type() const override { return "osd"; }
|
|
|
|
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
|
|
id_ = config.ValueOr<std::string>("id", "osd");
|
|
draw_bbox_ = config.ValueOr<bool>("draw_bbox", true);
|
|
draw_text_ = config.ValueOr<bool>("draw_text", true);
|
|
line_width_ = config.ValueOr<int>("line_width", 2);
|
|
font_scale_ = config.ValueOr<int>("font_scale", 1);
|
|
|
|
input_queue_ = ctx.input_queue;
|
|
if (!input_queue_) {
|
|
std::cerr << "[osd] no input queue for node " << id_ << "\n";
|
|
return false;
|
|
}
|
|
if (ctx.output_queues.empty()) {
|
|
std::cerr << "[osd] no output queue for node " << id_ << "\n";
|
|
return false;
|
|
}
|
|
output_queues_ = ctx.output_queues;
|
|
return true;
|
|
}
|
|
|
|
bool Start() override {
|
|
if (!input_queue_) return false;
|
|
running_.store(true);
|
|
worker_ = std::thread(&OsdNode::WorkerLoop, this);
|
|
std::cout << "[osd] started, draw_bbox=" << draw_bbox_ << " draw_text=" << draw_text_ << "\n";
|
|
return true;
|
|
}
|
|
|
|
void Stop() override {
|
|
running_.store(false);
|
|
if (input_queue_) input_queue_->Stop();
|
|
for (auto& q : output_queues_) q->Stop();
|
|
if (worker_.joinable()) worker_.join();
|
|
std::cout << "[osd] stopped\n";
|
|
}
|
|
|
|
private:
|
|
void PushToDownstream(FramePtr frame) {
|
|
for (auto& q : output_queues_) {
|
|
q->Push(frame);
|
|
}
|
|
}
|
|
|
|
void WorkerLoop() {
|
|
using namespace std::chrono;
|
|
FramePtr frame;
|
|
|
|
while (running_.load()) {
|
|
if (!input_queue_->Pop(frame, milliseconds(200))) continue;
|
|
if (!frame) continue;
|
|
|
|
if (frame->det && frame->data) {
|
|
DrawDetections(frame);
|
|
}
|
|
|
|
PushToDownstream(frame);
|
|
++processed_;
|
|
|
|
if (processed_ % 100 == 0) {
|
|
std::cout << "[osd] processed " << processed_ << " frames\n";
|
|
}
|
|
}
|
|
}
|
|
|
|
void DrawDetections(FramePtr frame) {
|
|
if (!frame->det || frame->det->items.empty()) return;
|
|
|
|
int w = frame->width;
|
|
int h = frame->height;
|
|
int stride = frame->stride > 0 ? frame->stride : w * 3;
|
|
uint8_t* data = frame->data;
|
|
PixelFormat fmt = frame->format;
|
|
|
|
bool supported = (fmt == PixelFormat::RGB || fmt == PixelFormat::BGR || fmt == PixelFormat::NV12);
|
|
if (!supported) {
|
|
return;
|
|
}
|
|
|
|
for (const auto& det : frame->det->items) {
|
|
int x1 = static_cast<int>(det.bbox.x);
|
|
int y1 = static_cast<int>(det.bbox.y);
|
|
int x2 = static_cast<int>(det.bbox.x + det.bbox.w);
|
|
int y2 = static_cast<int>(det.bbox.y + det.bbox.h);
|
|
|
|
Color color = GetClassColor(det.cls_id);
|
|
|
|
if (draw_bbox_) {
|
|
DrawRect(data, w, h, stride, fmt, x1, y1, x2, y2, line_width_, color);
|
|
}
|
|
|
|
if (draw_text_) {
|
|
char label[64];
|
|
snprintf(label, sizeof(label), "%s %.0f%%", GetClassName(det.cls_id), det.score * 100);
|
|
int text_y = y1 - 8 * font_scale_;
|
|
if (text_y < 0) text_y = y1 + 2;
|
|
DrawText(data, w, h, stride, fmt, x1, text_y, label, font_scale_, color);
|
|
}
|
|
}
|
|
}
|
|
|
|
std::string id_;
|
|
bool draw_bbox_ = true;
|
|
bool draw_text_ = true;
|
|
int line_width_ = 2;
|
|
int font_scale_ = 1;
|
|
|
|
std::atomic<bool> running_{false};
|
|
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
|
|
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
|
std::thread worker_;
|
|
uint64_t processed_ = 0;
|
|
};
|
|
|
|
REGISTER_NODE(OsdNode, "osd");
|
|
|
|
} // namespace rk3588
|