diff --git a/opensfm/src/map/pymap.pyi b/opensfm/src/map/pymap.pyi index f5f09dcf1..d6a5276b3 100644 --- a/opensfm/src/map/pymap.pyi +++ b/opensfm/src/map/pymap.pyi @@ -585,6 +585,8 @@ class TracksManager: def num_shots(self) -> int: ... def num_tracks(self) -> int: ... def remove_observation(self, arg0: str, arg1: str) -> None: ... + def remove_track(self, arg0: str) -> None: ... + def remove_shot(self, arg0: str) -> None: ... def write_to_file(self, arg0: str) -> None: ... Angular: "ErrorType" diff --git a/opensfm/src/map/python/pybind.cc b/opensfm/src/map/python/pybind.cc index fcc0209ec..cd047dbe3 100644 --- a/opensfm/src/map/python/pybind.cc +++ b/opensfm/src/map/python/pybind.cc @@ -292,6 +292,8 @@ PYBIND11_MODULE(pymap, m) { &map::TracksManager::MergeTracksManager) .def("add_observation", &map::TracksManager::AddObservation) .def("remove_observation", &map::TracksManager::RemoveObservation) + .def("remove_track", &map::TracksManager::RemoveTrack) + .def("remove_shot", &map::TracksManager::RemoveShot) .def("num_shots", &map::TracksManager::NumShots) .def("num_tracks", &map::TracksManager::NumTracks) .def("get_shot_ids", &map::TracksManager::GetShotIds) diff --git a/opensfm/src/map/src/tracks_manager.cc b/opensfm/src/map/src/tracks_manager.cc index 89dc64aa4..7174378f2 100644 --- a/opensfm/src/map/src/tracks_manager.cc +++ b/opensfm/src/map/src/tracks_manager.cc @@ -203,6 +203,46 @@ void TracksManager::RemoveObservation(const ShotId& shot_id, find_track->second.erase(shot_id); } +void TracksManager::RemoveTrack(const TrackId& track_id) { + // Step 1: Remove entries from tracks_per_shot_ + if (shots_per_track_.count(track_id)) { + const auto& shot_ids = shots_per_track_[track_id]; + for (const auto& [shot_id, _] : shot_ids) { + if (tracks_per_shot_.count(shot_id)) { + tracks_per_shot_[shot_id].erase(track_id); + // Remove the Shot entry if no observations remain + if (tracks_per_shot_[shot_id].empty()) { + tracks_per_shot_.erase(shot_id); + } + } + } + // Step 2: Remove the track from shots_per_track_ + shots_per_track_.erase(track_id); + } +} + +void TracksManager::RemoveShot(const ShotId& shot_id) { + // Remove shot from tracks_per_shot_ + const auto find_shot = tracks_per_shot_.find(shot_id); + if (find_shot != tracks_per_shot_.end()) { + // Iterate all tracks which are observed by the shot + for (const auto& [track_id, _] : find_shot->second) { + // Remove the shot from the corrsponding track-list in shots_per_track_ + if (shots_per_track_.count(track_id)) { + shots_per_track_[track_id].erase(shot_id); + // Remove the track entry if no observations remain + if (shots_per_track_[track_id].empty()) { + shots_per_track_.erase(track_id); + } + } + } + tracks_per_shot_.erase(shot_id); + } + else { + throw std::runtime_error("Shot ID not found"); + } +} + int TracksManager::NumShots() const { return tracks_per_shot_.size(); } int TracksManager::NumTracks() const { return shots_per_track_.size(); } diff --git a/opensfm/src/map/tracks_manager.h b/opensfm/src/map/tracks_manager.h index ab17d3189..6880bf7f3 100644 --- a/opensfm/src/map/tracks_manager.h +++ b/opensfm/src/map/tracks_manager.h @@ -14,6 +14,8 @@ class TracksManager { void AddObservation(const ShotId& shot_id, const TrackId& track_id, const Observation& observation); void RemoveObservation(const ShotId& shot_id, const TrackId& track_id); + void RemoveTrack(const TrackId& track_id); + void RemoveShot(const ShotId& shot_id); Observation GetObservation(const ShotId& shot, const TrackId& track) const; int NumShots() const;