118 lines
4.0 KiB
C++
118 lines
4.0 KiB
C++
#include <gtest/gtest.h>
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
|
|
#include "frame/frame.h"
|
|
#include "node.h"
|
|
#include "pose/pose_result.h"
|
|
#include "utils/simple_json.h"
|
|
#include "../plugins/pose_assoc/pose_assoc_node.h"
|
|
|
|
namespace rk3588 {
|
|
namespace {
|
|
|
|
SimpleJson ParsePoseAssocConfig(const std::string& text) {
|
|
SimpleJson config;
|
|
std::string err;
|
|
const bool ok = ParseSimpleJson(text, config, err);
|
|
EXPECT_TRUE(ok);
|
|
return config;
|
|
}
|
|
|
|
TEST(PoseAssocNodeTest, AssociatesPoseItemsToTrackedDetections) {
|
|
PoseAssocNode node;
|
|
SimpleJson config = ParsePoseAssocConfig(R"({
|
|
"id": "pose_assoc0",
|
|
"min_iou": 0.1
|
|
})");
|
|
|
|
NodeContext ctx;
|
|
auto out = std::make_shared<SpscQueue<FramePtr>>(4, QueueDropStrategy::DropOldest);
|
|
ctx.output_queues.push_back(out);
|
|
|
|
ASSERT_TRUE(node.Init(config, ctx));
|
|
ASSERT_TRUE(node.Start());
|
|
|
|
auto frame = std::make_shared<Frame>();
|
|
frame->det = std::make_shared<DetectionResult>();
|
|
frame->det->items.push_back(Detection{0, 0.95f, Rect{100.0f, 100.0f, 120.0f, 240.0f}, 7});
|
|
frame->det->items.push_back(Detection{0, 0.94f, Rect{320.0f, 120.0f, 120.0f, 240.0f}, 8});
|
|
|
|
frame->pose = std::make_shared<PoseResult>();
|
|
PoseItem pose0;
|
|
pose0.bbox = Rect{110.0f, 110.0f, 110.0f, 220.0f};
|
|
PoseItem pose1;
|
|
pose1.bbox = Rect{330.0f, 130.0f, 100.0f, 210.0f};
|
|
frame->pose->items.push_back(pose0);
|
|
frame->pose->items.push_back(pose1);
|
|
|
|
EXPECT_EQ(static_cast<int>(node.Process(frame)), static_cast<int>(NodeStatus::OK));
|
|
ASSERT_EQ(frame->pose->items.size(), 2u);
|
|
EXPECT_EQ(frame->pose->items[0].track_id, 7);
|
|
EXPECT_EQ(frame->pose->items[1].track_id, 8);
|
|
|
|
FramePtr forwarded;
|
|
ASSERT_TRUE(out->TryPop(forwarded));
|
|
EXPECT_EQ(forwarded.get(), frame.get());
|
|
}
|
|
|
|
TEST(PoseAssocNodeTest, UsesOneToOneAssignmentForOverlappingCandidates) {
|
|
PoseAssocNode node;
|
|
SimpleJson config = ParsePoseAssocConfig(R"({
|
|
"id": "pose_assoc0",
|
|
"min_iou": 0.05
|
|
})");
|
|
|
|
NodeContext ctx;
|
|
ASSERT_TRUE(node.Init(config, ctx));
|
|
ASSERT_TRUE(node.Start());
|
|
|
|
auto frame = std::make_shared<Frame>();
|
|
frame->det = std::make_shared<DetectionResult>();
|
|
frame->det->items.push_back(Detection{0, 0.95f, Rect{100.0f, 100.0f, 150.0f, 250.0f}, 11});
|
|
frame->det->items.push_back(Detection{0, 0.94f, Rect{130.0f, 120.0f, 150.0f, 250.0f}, 12});
|
|
|
|
frame->pose = std::make_shared<PoseResult>();
|
|
PoseItem pose0;
|
|
pose0.bbox = Rect{105.0f, 110.0f, 145.0f, 245.0f};
|
|
PoseItem pose1;
|
|
pose1.bbox = Rect{132.0f, 125.0f, 148.0f, 245.0f};
|
|
frame->pose->items.push_back(pose0);
|
|
frame->pose->items.push_back(pose1);
|
|
|
|
EXPECT_EQ(static_cast<int>(node.Process(frame)), static_cast<int>(NodeStatus::OK));
|
|
EXPECT_NE(frame->pose->items[0].track_id, frame->pose->items[1].track_id);
|
|
EXPECT_TRUE(frame->pose->items[0].track_id == 11 || frame->pose->items[0].track_id == 12);
|
|
EXPECT_TRUE(frame->pose->items[1].track_id == 11 || frame->pose->items[1].track_id == 12);
|
|
}
|
|
|
|
TEST(PoseAssocNodeTest, PassesThroughWhenPoseOrDetectionMissing) {
|
|
PoseAssocNode node;
|
|
SimpleJson config = ParsePoseAssocConfig(R"({
|
|
"id": "pose_assoc0"
|
|
})");
|
|
|
|
NodeContext ctx;
|
|
ASSERT_TRUE(node.Init(config, ctx));
|
|
ASSERT_TRUE(node.Start());
|
|
|
|
auto no_pose = std::make_shared<Frame>();
|
|
no_pose->det = std::make_shared<DetectionResult>();
|
|
no_pose->det->items.push_back(Detection{0, 0.9f, Rect{0.0f, 0.0f, 10.0f, 10.0f}, 1});
|
|
EXPECT_EQ(static_cast<int>(node.Process(no_pose)), static_cast<int>(NodeStatus::OK));
|
|
EXPECT_EQ(no_pose->pose, nullptr);
|
|
|
|
auto no_det = std::make_shared<Frame>();
|
|
no_det->pose = std::make_shared<PoseResult>();
|
|
PoseItem pose;
|
|
pose.bbox = Rect{0.0f, 0.0f, 10.0f, 10.0f};
|
|
no_det->pose->items.push_back(pose);
|
|
EXPECT_EQ(static_cast<int>(node.Process(no_det)), static_cast<int>(NodeStatus::OK));
|
|
ASSERT_EQ(no_det->pose->items.size(), 1u);
|
|
EXPECT_EQ(no_det->pose->items[0].track_id, -1);
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace rk3588
|