OrangePi3588Media/plugins/osd/osd_node.cpp
2025-12-31 15:56:40 +08:00

445 lines
17 KiB
C++

#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdio>
#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 || fmt == PixelFormat::YUV420) {
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 * stride + 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 || fmt == PixelFormat::YUV420) {
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 * stride + 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 || fmt == PixelFormat::YUV420) {
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
data[py * stride + 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);
// Load custom labels from config
if (const SimpleJson* labels = config.Find("labels")) {
for (const auto& item : labels->AsArray()) {
labels_.push_back(item.AsString(""));
}
std::cout << "[osd] loaded " << labels_.size() << " custom labels\n";
}
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 {
std::cout << "[osd] started, draw_bbox=" << draw_bbox_ << " draw_text=" << draw_text_ << "\n";
return true;
}
void Stop() override {
std::cout << "[osd] stopped\n";
}
NodeStatus Process(FramePtr frame) override {
if (!frame) return NodeStatus::DROP;
FramePtr out = frame;
if (out->det && out->data) {
if (out.use_count() > 1) {
out = CloneFrameForWrite(out);
}
DrawDetections(out);
}
PushToDownstream(out);
++processed_;
if (processed_ % 100 == 0) {
std::cout << "[osd] processed " << processed_ << " frames\n";
}
return NodeStatus::OK;
}
private:
static FramePtr CloneFrameForWrite(const FramePtr& src) {
if (!src || !src->data || src->data_size == 0) return src;
auto buf = std::make_shared<std::vector<uint8_t>>(src->data_size);
memcpy(buf->data(), src->data, src->data_size);
auto out = std::make_shared<Frame>(*src);
out->dma_fd = -1;
out->data = buf->data();
out->data_size = buf->size();
out->data_owner = buf;
for (int i = 0; i < 3; ++i) {
FramePlane p = src->planes[i];
int offset = p.offset;
if (p.data && src->data) {
const ptrdiff_t d = p.data - src->data;
if (d >= 0 && static_cast<size_t>(d) < src->data_size) {
offset = static_cast<int>(d);
}
}
p.offset = offset;
if (offset >= 0 && static_cast<size_t>(offset) < out->data_size) {
p.data = out->data + offset;
} else {
p.data = nullptr;
}
out->planes[i] = p;
}
return out;
}
void PushToDownstream(FramePtr frame) {
for (auto& q : output_queues_) {
q->Push(frame);
}
}
const char* GetLabel(int cls_id) const {
if (!labels_.empty()) {
if (cls_id >= 0 && cls_id < static_cast<int>(labels_.size())) {
return labels_[cls_id].c_str();
}
return "unknown";
}
return GetClassName(cls_id);
}
void DrawDetections(FramePtr frame) {
if (!frame->det || frame->det->items.empty()) return;
int w = frame->width;
int h = frame->height;
uint8_t* data = frame->planes[0].data ? frame->planes[0].data : frame->data;
PixelFormat fmt = frame->format;
int stride;
if (fmt == PixelFormat::RGB || fmt == PixelFormat::BGR) {
stride = frame->planes[0].stride > 0 ? frame->planes[0].stride : (frame->stride > 0 ? frame->stride : w * 3);
} else if (fmt == PixelFormat::NV12 || fmt == PixelFormat::YUV420) {
stride = frame->planes[0].stride > 0 ? frame->planes[0].stride : (frame->stride > 0 ? frame->stride : w);
} else {
return; // Unsupported format
}
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%%", GetLabel(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::vector<std::string> labels_;
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
uint64_t processed_ = 0;
};
REGISTER_NODE(OsdNode, "osd");
} // namespace rk3588