Add tracked triangulation association reports
This commit is contained in:
+77
-13
@@ -20,6 +20,7 @@ namespace
|
||||
using PoseArray2D =
|
||||
nb::ndarray<nb::numpy, const float, nb::shape<-1, -1, -1, 3>, nb::c_contig>;
|
||||
using CountArray = nb::ndarray<nb::numpy, const uint32_t, nb::shape<-1>, nb::c_contig>;
|
||||
using TrackIdArray = nb::ndarray<nb::numpy, const int64_t, nb::shape<-1>, nb::c_contig>;
|
||||
using PoseArray3DConst =
|
||||
nb::ndarray<nb::numpy, const float, nb::shape<-1, -1, 4>, nb::c_contig>;
|
||||
using PoseArray3D = nb::ndarray<nb::numpy, float, nb::shape<-1, -1, 4>, nb::c_contig>;
|
||||
@@ -58,6 +59,24 @@ PoseBatch3DView pose_batch3d_view_from_numpy(const PoseArray3DConst &poses_3d)
|
||||
};
|
||||
}
|
||||
|
||||
TrackedPoseBatch3DView tracked_pose_batch_view_from_numpy(
|
||||
const PoseArray3DConst &poses_3d,
|
||||
const TrackIdArray &track_ids)
|
||||
{
|
||||
if (poses_3d.shape(0) != track_ids.shape(0))
|
||||
{
|
||||
throw std::invalid_argument(
|
||||
"previous_poses_3d and previous_track_ids must have the same number of tracks.");
|
||||
}
|
||||
|
||||
return TrackedPoseBatch3DView {
|
||||
track_ids.data(),
|
||||
poses_3d.data(),
|
||||
static_cast<size_t>(poses_3d.shape(0)),
|
||||
static_cast<size_t>(poses_3d.shape(1)),
|
||||
};
|
||||
}
|
||||
|
||||
PoseArray3D pose_batch_to_numpy(PoseBatch3D batch)
|
||||
{
|
||||
auto *storage = new std::vector<float>(std::move(batch.data));
|
||||
@@ -216,6 +235,7 @@ NB_MODULE(_core, m)
|
||||
nb::class_<PreviousPoseMatch>(m, "PreviousPoseMatch")
|
||||
.def(nb::init<>())
|
||||
.def_rw("previous_pose_index", &PreviousPoseMatch::previous_pose_index)
|
||||
.def_rw("previous_track_id", &PreviousPoseMatch::previous_track_id)
|
||||
.def_rw("score_view1", &PreviousPoseMatch::score_view1)
|
||||
.def_rw("score_view2", &PreviousPoseMatch::score_view2)
|
||||
.def_rw("matched_view1", &PreviousPoseMatch::matched_view1)
|
||||
@@ -274,6 +294,34 @@ NB_MODULE(_core, m)
|
||||
return merged_poses_to_numpy_copy(merge.merged_poses);
|
||||
}, nb::rv_policy::move);
|
||||
|
||||
nb::enum_<AssociationStatus>(m, "AssociationStatus")
|
||||
.value("MATCHED", AssociationStatus::Matched)
|
||||
.value("NEW", AssociationStatus::New)
|
||||
.value("AMBIGUOUS", AssociationStatus::Ambiguous);
|
||||
|
||||
nb::class_<AssociationReport>(m, "AssociationReport")
|
||||
.def(nb::init<>())
|
||||
.def_rw("pose_previous_indices", &AssociationReport::pose_previous_indices)
|
||||
.def_rw("pose_previous_track_ids", &AssociationReport::pose_previous_track_ids)
|
||||
.def_rw("pose_status", &AssociationReport::pose_status)
|
||||
.def_rw("pose_candidate_previous_indices", &AssociationReport::pose_candidate_previous_indices)
|
||||
.def_rw("pose_candidate_previous_track_ids", &AssociationReport::pose_candidate_previous_track_ids)
|
||||
.def_rw("unmatched_previous_indices", &AssociationReport::unmatched_previous_indices)
|
||||
.def_rw("unmatched_previous_track_ids", &AssociationReport::unmatched_previous_track_ids)
|
||||
.def_rw("new_pose_indices", &AssociationReport::new_pose_indices)
|
||||
.def_rw("ambiguous_pose_indices", &AssociationReport::ambiguous_pose_indices);
|
||||
|
||||
nb::class_<FinalPoseAssociationDebug>(m, "FinalPoseAssociationDebug")
|
||||
.def(nb::init<>())
|
||||
.def_rw("final_pose_index", &FinalPoseAssociationDebug::final_pose_index)
|
||||
.def_rw("source_core_proposal_indices", &FinalPoseAssociationDebug::source_core_proposal_indices)
|
||||
.def_rw("source_pair_indices", &FinalPoseAssociationDebug::source_pair_indices)
|
||||
.def_rw("candidate_previous_indices", &FinalPoseAssociationDebug::candidate_previous_indices)
|
||||
.def_rw("candidate_previous_track_ids", &FinalPoseAssociationDebug::candidate_previous_track_ids)
|
||||
.def_rw("resolved_previous_index", &FinalPoseAssociationDebug::resolved_previous_index)
|
||||
.def_rw("resolved_previous_track_id", &FinalPoseAssociationDebug::resolved_previous_track_id)
|
||||
.def_rw("status", &FinalPoseAssociationDebug::status);
|
||||
|
||||
nb::class_<TriangulationTrace>(m, "TriangulationTrace")
|
||||
.def(nb::init<>())
|
||||
.def_rw("pairs", &TriangulationTrace::pairs)
|
||||
@@ -282,11 +330,21 @@ NB_MODULE(_core, m)
|
||||
.def_rw("grouping", &TriangulationTrace::grouping)
|
||||
.def_rw("full_proposals", &TriangulationTrace::full_proposals)
|
||||
.def_rw("merge", &TriangulationTrace::merge)
|
||||
.def_rw("association", &TriangulationTrace::association)
|
||||
.def_rw("final_pose_associations", &TriangulationTrace::final_pose_associations)
|
||||
.def_prop_ro("final_poses", [](const TriangulationTrace &trace)
|
||||
{
|
||||
return pose_batch_to_numpy_copy(trace.final_poses);
|
||||
}, nb::rv_policy::move);
|
||||
|
||||
nb::class_<TriangulationResult>(m, "TriangulationResult")
|
||||
.def(nb::init<>())
|
||||
.def_rw("association", &TriangulationResult::association)
|
||||
.def_prop_ro("poses_3d", [](const TriangulationResult &result)
|
||||
{
|
||||
return pose_batch_to_numpy_copy(result.poses);
|
||||
}, nb::rv_policy::move);
|
||||
|
||||
m.def(
|
||||
"make_camera",
|
||||
&make_camera,
|
||||
@@ -313,17 +371,19 @@ NB_MODULE(_core, m)
|
||||
[](const PoseArray2D &poses_2d,
|
||||
const CountArray &person_counts,
|
||||
const TriangulationConfig &config,
|
||||
const PoseArray3DConst &previous_poses_3d)
|
||||
const PoseArray3DConst &previous_poses_3d,
|
||||
const TrackIdArray &previous_track_ids)
|
||||
{
|
||||
return filter_pairs_with_previous_poses(
|
||||
pose_batch_view_from_numpy(poses_2d, person_counts),
|
||||
config,
|
||||
pose_batch3d_view_from_numpy(previous_poses_3d));
|
||||
tracked_pose_batch_view_from_numpy(previous_poses_3d, previous_track_ids));
|
||||
},
|
||||
"poses_2d"_a,
|
||||
"person_counts"_a,
|
||||
"config"_a,
|
||||
"previous_poses_3d"_a);
|
||||
"previous_poses_3d"_a,
|
||||
"previous_track_ids"_a);
|
||||
|
||||
m.def(
|
||||
"triangulate_debug",
|
||||
@@ -342,9 +402,11 @@ NB_MODULE(_core, m)
|
||||
[](const PoseArray2D &poses_2d,
|
||||
const CountArray &person_counts,
|
||||
const TriangulationConfig &config,
|
||||
const PoseArray3DConst &previous_poses_3d)
|
||||
const PoseArray3DConst &previous_poses_3d,
|
||||
const TrackIdArray &previous_track_ids)
|
||||
{
|
||||
const PoseBatch3DView previous_view = pose_batch3d_view_from_numpy(previous_poses_3d);
|
||||
const TrackedPoseBatch3DView previous_view =
|
||||
tracked_pose_batch_view_from_numpy(previous_poses_3d, previous_track_ids);
|
||||
return triangulate_debug(
|
||||
pose_batch_view_from_numpy(poses_2d, person_counts),
|
||||
config,
|
||||
@@ -353,7 +415,8 @@ NB_MODULE(_core, m)
|
||||
"poses_2d"_a,
|
||||
"person_counts"_a,
|
||||
"config"_a,
|
||||
"previous_poses_3d"_a);
|
||||
"previous_poses_3d"_a,
|
||||
"previous_track_ids"_a);
|
||||
|
||||
m.def(
|
||||
"triangulate_poses",
|
||||
@@ -370,21 +433,22 @@ NB_MODULE(_core, m)
|
||||
"config"_a);
|
||||
|
||||
m.def(
|
||||
"triangulate_poses",
|
||||
"triangulate_with_report",
|
||||
[](const PoseArray2D &poses_2d,
|
||||
const CountArray &person_counts,
|
||||
const TriangulationConfig &config,
|
||||
const PoseArray3DConst &previous_poses_3d)
|
||||
const PoseArray3DConst &previous_poses_3d,
|
||||
const TrackIdArray &previous_track_ids)
|
||||
{
|
||||
const PoseBatch3DView previous_view = pose_batch3d_view_from_numpy(previous_poses_3d);
|
||||
const PoseBatch3D poses_3d = triangulate_poses(
|
||||
const TriangulationResult result = triangulate_with_report(
|
||||
pose_batch_view_from_numpy(poses_2d, person_counts),
|
||||
config,
|
||||
&previous_view);
|
||||
return pose_batch_to_numpy(poses_3d);
|
||||
tracked_pose_batch_view_from_numpy(previous_poses_3d, previous_track_ids));
|
||||
return result;
|
||||
},
|
||||
"poses_2d"_a,
|
||||
"person_counts"_a,
|
||||
"config"_a,
|
||||
"previous_poses_3d"_a);
|
||||
"previous_poses_3d"_a,
|
||||
"previous_track_ids"_a);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user