safesight-edge/plugins/preprocess/preprocess_node.cpp

287 lines
9.9 KiB
C++

#include <algorithm>
#include <cstdint>
#include <string>
#include <vector>
#include "face/face_result.h"
#include "hw/i_image_processor.h"
#include "node.h"
#include "utils/logger.h"
namespace rk3588 {
namespace {
PixelFormat ParseFormat(const std::string& s) {
if (s == "nv12" || s == "NV12") return PixelFormat::NV12;
if (s == "yuv420" || s == "YUV420") return PixelFormat::YUV420;
if (s == "rgb" || s == "RGB") return PixelFormat::RGB;
if (s == "bgr" || s == "BGR") return PixelFormat::BGR;
return PixelFormat::UNKNOWN;
}
inline float ClampFloat(float v, float lo, float hi) {
return std::max(lo, std::min(v, hi));
}
void ScaleRect(Rect& r, float sx, float sy, int out_w, int out_h) {
if (out_w <= 0 || out_h <= 0) {
r = Rect{};
return;
}
const float fw = static_cast<float>(out_w);
const float fh = static_cast<float>(out_h);
const float x = ClampFloat(r.x * sx, 0.0f, fw);
const float y = ClampFloat(r.y * sy, 0.0f, fh);
float w = std::max(0.0f, r.w * sx);
float h = std::max(0.0f, r.h * sy);
if (x + w > fw) w = std::max(0.0f, fw - x);
if (y + h > fh) h = std::max(0.0f, fh - y);
r.x = x;
r.y = y;
r.w = w;
r.h = h;
}
void ScalePoint(Point2f& p, float sx, float sy, int out_w, int out_h) {
if (out_w <= 0 || out_h <= 0) {
p = Point2f{};
return;
}
p.x = ClampFloat(p.x * sx, 0.0f, static_cast<float>(out_w));
p.y = ClampFloat(p.y * sy, 0.0f, static_cast<float>(out_h));
}
} // namespace
class PreprocessNode : public INode {
public:
std::string Id() const override { return id_; }
std::string Type() const override { return "preprocess"; }
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
id_ = config.ValueOr<std::string>("id", "preprocess");
dst_w_ = config.ValueOr<int>("dst_w", 640);
dst_h_ = config.ValueOr<int>("dst_h", 640);
keep_ratio_ = config.ValueOr<bool>("keep_ratio", false);
std::string fmt_str = config.ValueOr<std::string>("dst_format", "");
if (!fmt_str.empty()) {
dst_fmt_ = ParseFormat(fmt_str);
}
const bool requested_use_rga = config.ValueOr<bool>("use_rga", true);
use_rga_ = requested_use_rga;
if (const SimpleJson* dbg = config.Find("debug"); dbg && dbg->IsObject()) {
stats_log_ = dbg->ValueOr<bool>("stats", stats_log_);
stats_interval_ = std::max<uint64_t>(
1, static_cast<uint64_t>(dbg->ValueOr<int>("stats_interval", static_cast<int>(stats_interval_))));
}
input_queue_ = ctx.input_queue;
if (!input_queue_) {
LogError("[preprocess] no input queue for node " + id_);
return false;
}
if (ctx.output_queues.empty()) {
LogError("[preprocess] no output queue for node " + id_);
return false;
}
output_queues_ = ctx.output_queues;
#if !defined(RK3588_ENABLE_RGA)
if (requested_use_rga) {
LogError("[preprocess] use_rga=true but RGA not enabled at build time");
return false;
}
use_rga_ = false;
#endif
#if !defined(RK3588_ENABLE_FFMPEG)
if (!use_rga_) {
LogError("[preprocess] neither RGA nor FFmpeg enabled");
return false;
}
#endif
image_processor_ = ctx.image_processor;
if (!image_processor_) {
LogError("[preprocess] no image processor for node " + id_);
return false;
}
return true;
}
bool Start() override {
LogInfo("[preprocess] start id=" + id_ + " dst=" + std::to_string(dst_w_) + "x" +
std::to_string(dst_h_) + (use_rga_ ? " (rga)" : " (swscale)"));
return true;
}
void Stop() override {}
NodeStatus Process(FramePtr frame) override {
if (!frame) return NodeStatus::DROP;
if (!image_processor_) return NodeStatus::ERROR;
PixelFormat out_fmt = (dst_fmt_ != PixelFormat::UNKNOWN) ? dst_fmt_ : frame->format;
int out_w = dst_w_;
int out_h = dst_h_;
if (out_w <= 0) out_w = frame->width;
if (out_h <= 0) out_h = frame->height;
if (keep_ratio_ && dst_w_ > 0 && dst_h_ > 0 && frame->width > 0 && frame->height > 0) {
float scale = std::min(static_cast<float>(dst_w_) / frame->width,
static_cast<float>(dst_h_) / frame->height);
out_w = static_cast<int>(frame->width * scale);
out_h = static_cast<int>(frame->height * scale);
out_w = (out_w + 1) & ~1;
out_h = (out_h + 1) & ~1;
}
const bool need_resize = (frame->width != out_w || frame->height != out_h);
const bool need_cvt = (frame->format != out_fmt);
if (need_resize) {
WarnMetaResizeOnce(frame, out_w, out_h);
}
if (!need_resize && !need_cvt) {
ProcessPassthrough(frame);
return NodeStatus::OK;
}
Frame out;
out.width = out_w;
out.height = out_h;
out.format = out_fmt;
Status st = image_processor_->Resize(*frame, out);
if (st.Failed()) {
if (!use_rga_ && st.ErrMessage().find("unsupported format") != std::string::npos) {
ProcessPassthrough(frame);
return NodeStatus::OK;
}
LogError("[preprocess] " + st.ErrMessage());
return NodeStatus::ERROR;
}
auto out_frame = std::make_shared<Frame>(out);
out_frame->pts = frame->pts;
out_frame->frame_id = frame->frame_id;
if (frame->det) {
auto det = std::make_shared<DetectionResult>(*frame->det);
if (frame->width > 0 && frame->height > 0) {
const float sx = static_cast<float>(out_w) / static_cast<float>(frame->width);
const float sy = static_cast<float>(out_h) / static_cast<float>(frame->height);
for (auto& it : det->items) {
ScaleRect(it.bbox, sx, sy, out_w, out_h);
}
}
det->img_w = out_w;
det->img_h = out_h;
out_frame->det = std::move(det);
}
if (frame->face_det) {
auto face_det = std::make_shared<FaceDetResult>(*frame->face_det);
if (frame->width > 0 && frame->height > 0) {
const float sx = static_cast<float>(out_w) / static_cast<float>(frame->width);
const float sy = static_cast<float>(out_h) / static_cast<float>(frame->height);
for (auto& it : face_det->faces) {
ScaleRect(it.bbox, sx, sy, out_w, out_h);
if (it.has_landmarks) {
for (auto& lm : it.landmarks) {
ScalePoint(lm, sx, sy, out_w, out_h);
}
}
}
}
face_det->img_w = out_w;
face_det->img_h = out_h;
out_frame->face_det = std::move(face_det);
}
if (frame->face_recog) {
auto face_recog = std::make_shared<FaceRecogResult>(*frame->face_recog);
if (frame->width > 0 && frame->height > 0) {
const float sx = static_cast<float>(out_w) / static_cast<float>(frame->width);
const float sy = static_cast<float>(out_h) / static_cast<float>(frame->height);
for (auto& it : face_recog->items) {
ScaleRect(it.bbox, sx, sy, out_w, out_h);
if (it.has_landmarks) {
for (auto& lm : it.landmarks) {
ScalePoint(lm, sx, sy, out_w, out_h);
}
}
}
}
face_recog->img_w = out_w;
face_recog->img_h = out_h;
out_frame->face_recog = std::move(face_recog);
}
out_frame->user_meta = frame->user_meta;
PushToDownstream(out_frame);
++processed_;
if (stats_log_ && stats_interval_ > 0 && (processed_ % stats_interval_) == 0) {
LogInfo("[preprocess] " + std::string(use_rga_ ? "rga" : "swscale") +
" frame=" + std::to_string(out_frame->frame_id) +
" " + std::to_string(frame->width) + "x" + std::to_string(frame->height) +
" -> " + std::to_string(out_w) + "x" + std::to_string(out_h) +
" id=" + id_);
}
return NodeStatus::OK;
}
private:
void PushToDownstream(FramePtr frame) {
for (auto& q : output_queues_) {
q->Push(frame);
}
}
void WarnMetaResizeOnce(const FramePtr& frame, int out_w, int out_h) {
if (warned_meta_resize_) return;
if (!frame) return;
if (frame->width == out_w && frame->height == out_h) return;
if (!frame->det && !frame->face_det && !frame->face_recog) return;
warned_meta_resize_ = true;
LogInfo("[preprocess] resized frame and scaled det/face meta to destination resolution (id=" + id_ + ")");
}
void ProcessPassthrough(FramePtr frame) {
PushToDownstream(frame);
++processed_;
if (stats_log_ && stats_interval_ > 0 && (processed_ % stats_interval_) == 0) {
LogInfo("[preprocess] passthrough frame=" + std::to_string(frame->frame_id) + " id=" + id_);
}
}
std::string id_;
int dst_w_ = 640;
int dst_h_ = 640;
bool keep_ratio_ = false;
PixelFormat dst_fmt_ = PixelFormat::UNKNOWN;
bool use_rga_ = true;
bool stats_log_ = false;
uint64_t stats_interval_ = 100;
bool warned_meta_resize_ = false;
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
std::shared_ptr<IImageProcessor> image_processor_;
uint64_t processed_ = 0;
};
REGISTER_NODE(PreprocessNode, "preprocess");
} // namespace rk3588