90 lines
3.4 KiB
C++
90 lines
3.4 KiB
C++
// Copyright 2019 The MediaPipe Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#ifndef MEDIAPIPE_UTIL_TRACKING_DETECTION_MANAGER_H_
|
|
#define MEDIAPIPE_UTIL_TRACKING_DETECTION_MANAGER_H_
|
|
|
|
#include "absl/container/node_hash_map.h"
|
|
#include "mediapipe/framework/formats/rect.pb.h"
|
|
#include "mediapipe/util/tracking/tracked_detection.h"
|
|
#include "mediapipe/util/tracking/tracked_detection_manager_config.pb.h"
|
|
|
|
namespace mediapipe {
|
|
|
|
// Class for managing detections.
|
|
// The detections are from either server side or a mobile detector and are
|
|
// tracked using either 2D or 3D tracker. The TrackedDetectionManager is used to
|
|
// identify duplicated detections or obsolete detections to keep a set of
|
|
// active detections.
|
|
class TrackedDetectionManager {
|
|
public:
|
|
TrackedDetectionManager() = default;
|
|
~TrackedDetectionManager() = default;
|
|
|
|
// Adds a new detection at a given timestamp. Returns the IDs of the
|
|
// detections that are removed due to duplication.
|
|
std::vector<int> AddDetection(std::unique_ptr<TrackedDetection> detection);
|
|
|
|
// Updates the location of a detection identified by its id.
|
|
// Returns false if the requested id doesn't exist a corresponding
|
|
// detection. Returns the IDs of the detections that are removed due to
|
|
// duplication.
|
|
std::vector<int> UpdateDetectionLocation(
|
|
int id, const ::mediapipe::NormalizedRect& bounding_box, int64 timestamp);
|
|
|
|
// Removes detections that are not updated after |timestamp|. Returns the IDs
|
|
// of the detections that are removed.
|
|
std::vector<int> RemoveObsoleteDetections(int64 timestamp);
|
|
|
|
// TODO: Do we really need this? Pursuit tracker doesn't do well
|
|
// in loop closure. Boxes out of view are usually not attached to objects
|
|
// any more when it's in view again. Returns the IDs of the
|
|
// detections that are removed.
|
|
std::vector<int> RemoveOutOfViewDetections();
|
|
|
|
int GetNumDetections() const { return detections_.size(); }
|
|
|
|
// Get TrackedDetection by its unique id.
|
|
const TrackedDetection* GetTrackedDetection(int id) const {
|
|
auto detection_ptr = detections_.find(id);
|
|
if (detection_ptr == detections_.end()) {
|
|
return nullptr;
|
|
}
|
|
return detection_ptr->second.get();
|
|
}
|
|
|
|
const absl::node_hash_map<int, std::unique_ptr<TrackedDetection>>&
|
|
GetAllTrackedDetections() const {
|
|
return detections_;
|
|
}
|
|
|
|
void SetConfig(const mediapipe::TrackedDetectionManagerConfig& config) {
|
|
config_ = config;
|
|
}
|
|
|
|
private:
|
|
// Finds all detections that are duplicated with the one of |id| and remove
|
|
// all detections except the one that is added most recently. Returns the IDs
|
|
// of the detections that are removed.
|
|
std::vector<int> RemoveDuplicatedDetections(int id);
|
|
|
|
absl::node_hash_map<int, std::unique_ptr<TrackedDetection>> detections_;
|
|
|
|
mediapipe::TrackedDetectionManagerConfig config_;
|
|
};
|
|
|
|
} // namespace mediapipe
|
|
|
|
#endif // MEDIAPIPE_UTIL_TRACKING_DETECTION_MANAGER_H_
|