talmolab / sleap

A deep learning framework for multi-animal pose tracking.
https://sleap.ai
Other
435 stars 97 forks source link

Refactor the max_tracking code #1896

Open getzze opened 3 months ago

getzze commented 3 months ago

ONLY AFTER #1894 is merged.

I refactored the max_tracking feature introduced in #1447 and simplified the implementation. It may give different results as the goal is similar but implemented differently.

One problem I had with the current implementation is that instances cannot be assigned to created but unused tracks. Say we have max_tracks=2, and the 2 tracks are created. If a track is not assigned instances for track_window timesteps, it does not contribute to instance candidates anymore, hence instances cannot be assigned to this track anymore, forever. What was happening is that instances stopped being tracked in the middle of the video. Therefore it was unusable for me and the "pre_cull" options are not available from the GUI anymore, so I need to use the CLI.

I added an option to reassign unused tracks to unmatched instances if the maximum number of tracks was reached.

I also added an option to merge instances in the "pre_cull" stage. That is because I often have two instances from inference: front of mouse and back of mouse, that could perfectly be merged. The idea is that only "good" merges are kept by the tracking step, but it is not perfect.

Summary by CodeRabbit

coderabbitai[bot] commented 3 months ago

Walkthrough

The recent updates enhance the tracking capabilities and usability of the system. Key modifications include streamlined command-line usage, refined configuration parameters, improved inference logic, and better handling of tracking instances. These changes collectively aim to simplify user interactions, optimize performance, and clarify documentation, ensuring a more robust and efficient tracking experience.

Changes

Files Change Summary
docs/guides/cli.md Updated command-line examples, replacing --max_tracking with --max_tracks to simplify options.
sleap/config/pipeline_form.yaml Removed tracking.max_tracking and added tracking.save_shifted_instances for improved tracking configuration.
sleap/gui/learning/runners.py Improved handling of inference parameters and tracker name compatibility.
sleap/instance.py Introduced all_disjoint and create_merged_instances functions for better instance processing.
sleap/nn/inference.py Enhanced prediction efficiency and backward compatibility, refined data processing workflow.
sleap/nn/tracker/components.py Added merging logic to cull_frame_instances for better instance management.
sleap/nn/tracking.py Restructured candidate-making logic with new abstract base class for flexibility.
tests/nn/test_inference.py Updated tests to reflect new tracking approach and improved performance handling.
tests/nn/test_tracker_components.py Renamed functions and added parameters for dynamic image scaling in tracking tests.
tests/nn/test_tracking_integration.py Refactored tracking tests for improved encapsulation and clarity.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant CLI
    participant Tracker
    participant Inference

    User->>CLI: Execute tracking command
    CLI->>Tracker: Initialize tracker
    Tracker->>Inference: Process data
    Inference-->>Tracker: Return results
    Tracker-->>CLI: Output results
    CLI-->>User: Display tracking results

πŸ‡ In the meadow, I hop with glee,
New changes made, oh what a spree!
Tracks refined and paths made clear,
With every hop, I shed a cheer!
Let’s bound ahead, no need to look back,
In our swift dance, we’re on the right track! 🌼✨


Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share - [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai) - [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai) - [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai) - [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)
Tips ### Chat There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai): - Review comments: Directly reply to a review comment made by CodeRabbit. Example: - `I pushed a fix in commit .` - `Generate unit testing code for this file.` - `Open a follow-up GitHub issue for this discussion.` - Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples: - `@coderabbitai generate unit testing code for this file.` - `@coderabbitai modularize this function.` - PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples: - `@coderabbitai generate interesting stats about this repository and render them as a table.` - `@coderabbitai show all the console.log statements in this repository.` - `@coderabbitai read src/utils.ts and generate unit testing code.` - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.` - `@coderabbitai help me debug CodeRabbit configuration file.` Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. ### CodeRabbit Commands (Invoked using PR comments) - `@coderabbitai pause` to pause the reviews on a PR. - `@coderabbitai resume` to resume the paused reviews. - `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository. - `@coderabbitai full review` to do a full review from scratch and review all the files again. - `@coderabbitai summary` to regenerate the summary of the PR. - `@coderabbitai resolve` resolve all the CodeRabbit review comments. - `@coderabbitai configuration` to show the current CodeRabbit configuration for the repository. - `@coderabbitai help` to get help. ### Other keywords and placeholders - Add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed. - Add `@coderabbitai summary` to generate the high-level summary at a specific location in the PR description. - Add `@coderabbitai` anywhere in the PR title to generate the title automatically. ### CodeRabbit Configuration File (`.coderabbit.yaml`) - You can programmatically configure CodeRabbit by adding a `.coderabbit.yaml` file to the root of your repository. - Please see the [configuration documentation](https://docs.coderabbit.ai/guides/configure-coderabbit) for more information. - If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: `# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json` ### Documentation and Community - Visit our [Documentation](https://coderabbit.ai/docs) for detailed information on how to use CodeRabbit. - Join our [Discord Community](https://discord.com/invite/GsXnASn26c) to get help, request features, and share feedback. - Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.
codecov[bot] commented 2 months ago

Codecov Report

Attention: Patch coverage is 65.46053% with 105 lines in your changes missing coverage. Please review.

Project coverage is 75.21%. Comparing base (7ed1229) to head (06ef136). Report is 56 commits behind head on develop.

Files with missing lines Patch % Lines
sleap/nn/tracking.py 74.28% 45 Missing :warning:
sleap/instance.py 15.62% 27 Missing :warning:
sleap/nn/inference.py 66.17% 23 Missing :warning:
sleap/nn/tracker/components.py 16.66% 5 Missing :warning:
sleap/gui/learning/runners.py 60.00% 2 Missing :warning:
sleap/gui/learning/dialog.py 0.00% 1 Missing :warning:
sleap/gui/widgets/video.py 50.00% 1 Missing :warning:
sleap/util.py 93.33% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## develop #1896 +/- ## =========================================== + Coverage 73.30% 75.21% +1.90% =========================================== Files 134 133 -1 Lines 24087 24558 +471 =========================================== + Hits 17658 18471 +813 + Misses 6429 6087 -342 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

gitttt-1234 commented 2 months ago

Hi @getzze,

Thanks for this!

We notice that you tried to simplify our tracking method by removing max_tracking. In SLEAP, we currently support 2 candidate generation methods: Fixed window and Local queues. In Fixed window method, instances from the last window_size frames are stored in the tracker queue. This method focuses on recent predictions but can lose a track if no detections are assigned for window_size frames. In Local queues method, the last window_size instances for each track are stored in the tracker queue. This method is better at connecting detections over gaps but doesn't rely on recent predictions. To support both the methods, we use two different data structures (a deque for fixed window method and a dict[track_id, deque] for local queues) and also max_tracking parameter to determine the type of candidate generation method to be used.

The PR refactors the code to implement a fixed window method with the option of limiting the number of new tracks being created. However, this does not support the local queues method, which is preferred in some cases (videos where tracks are to be connected after short gaps).

Let us know if you have any questions!

Thanks!

Divya

getzze commented 2 months ago

Hi @gitttt-1234,

Thanks for reviewing the PR!

I chose to remove the Local queues alternative because I thought its only point was to limit the number of tracks/identities. I didn't get that it could be used if only one instance is hidden for some time. Anyway, I think it's very unlikely that the instance would reappear exactly at the exact same place it disappeared, so it is matched to the track that disappeared. In my tests, I always got tracks that become empty after the instance is lost for window_size time steps, they are never "caught" again by the track.

Instead I added an (optional) step to assign an instance to a track that has disappeared for more than window_size time steps, which works better in my tests (two mice).

Something I didn't do in this PR and that should probably be done, is to put back the pre_cull_to_target option in the GUI, to cap the number of instances before assigning them to tracks (whereas max_tracks cap them after assigning them to tracks), to be able to get back the pre-#1447 behavior.

talmo commented 2 months ago

Hi @getzze,

Thanks for the contributions! Apologies it's taken us so long to work through it all, but we did have some questions and concerns about the changes. @gitttt-1234 described those inline, but I think maybe there's also some misunderstandings that I want to clarify to make sure we're on the same page -- see below.

Hi @gitttt-1234,

Thanks for reviewing the PR!

I chose to remove the Local queues alternative because I thought its only point was to limit the number of tracks/identities.

As @gitttt-1234 mentioned, the point of the local queues is indeed to limit the maximum number of tracks that can be possibly spawned in an entire project, but it does so using a specific algorithmic approach.

By specifying the maximum number of tracks, what we can do is create buffers (up to that limit) that hold the past window_size detected instances that have been assigned to that track. This means that if we don't see an instance from track_0 for 1000+ frames, we're still able to assign a new instance to track_0 back to that track.

The main thing this solves is that it prevents the explosion of new tracklets being created in scenarios where there is a lot of track fragmentation due to frequent occlusions (or lots of animals in the FOV).

I didn't get that it could be used if only one instance is hidden for some time.

Again, to be clear, with local track queues, the instances can be used no matter how many instances are hidden and no matter the amount of time.

Anyway, I think it's very unlikely that the instance would reappear exactly at the exact same place it disappeared, so it is matched to the track that disappeared.

Using instances from a long time in the past is not great since the detections will be pretty stale, but it can work well when animals are sleeping or hiding under a hut or something else out of view. It does not, however, mean that the instance needs to reappear in the exact same place. It just means that the last detection(s) from that track will still be in consideration during the candidate generation and matching steps.

If we lose a track for a long amount of time, the truly appropriate thing to do would be to explicitly use ReID features, but we don't have that implemented in here yet, so in the meantime, this is a "best effort" approach that doesn't require any new information to be integrated in the algorithm.

In my tests, I always got tracks that become empty after the instance is lost for window_size time steps, they are never "caught" again by the track.

Could you elaborate on what you mean by that? How does a track "become empty"?

Again, by design, if an instance is lost for more than window_size time steps, when using the local queues in simplemaxtracks or flowmaxtracks, last few instances that have been assigned to a track can never be lost no matter the gap length since they're stored in queues local to the track (ergo, "track local queues").

Instead I added an (optional) step to assign an instance to a track that has disappeared for more than window_size time steps, which works better in my tests (two mice).

The track local queues should handle this explicitly. We don't make a distinction between tracks that are "stale" vs "active", but treat all instances from all queues as valid candidates. I could see an advantage to separating these two (e.g., for preferring active tracks that have recent detections before considering the stale ones), but this is fundamentally a different tracking algorithm and not a replacement for what we have now.

Something I didn't do in this PR and that should probably be done, is to put back the pre_cull_to_target option in the GUI, to cap the number of instances before assigning them to tracks (whereas max_tracks cap them after assigning them to tracks), to be able to get back the pre-#1447 behavior.

We removed pre-culling as this was superseded by a feature set for limiting the number of instances right at the model inference level. This is exposed via the --max_instances flag.

We implemented these in:

The advantage is that it allows the users to cull to a specific number of instances regardless of whether tracking will be applied next -- so it's usable during the human-in-the-loop annotation stage as well to filter out some false positives when models aren't trained well enough yet.


We also had a question about the instance merging functionality you added. It makes total sense, but you might imagine how this will have lots of edge cases.

For example: if the two node sets are exactly disjoint, but correctly belong to two different animals, we'd be incorrectly merging them into what we call a "chimera" (top of one mouse connected to the bottom of another mouse).

Another scenario we see often that this wouldn't quite resolve is when we have two almost disjoint sets, except for one node that is maybe off by a couple of pixels. Adding a bit of fuzziness to the keypoint matching to test for "disjoint-ness" before triggering this function could help cover those cases as well.

The thing I'm a bit more confused by is that you're mentioning this as useful when there are two instances on the same frame that also belong to the same track. How is this even possible given any of the current tracking approaches? We always do matching that ensures that no more than one instance is assigned to a given track within a frame.

Establishing that two (partial) instances indeed belong to the same track is the hard problem that's at the crux of this instance merging problem -- and why we haven't implemented this before.

Some clarity on what you had in mind for dealing with this would help us understand the solution.


In general, we're super appreciative of the contributions, but we want to make sure that:

  1. New algorithms are applicable to the broader use cases that we see across the user base;
  2. We are very careful not to break backwards compatibility, especially when removing core algorithms; and
  3. We make best efforts to separate different changes into different PRs for future maintainability and being able to track down when bugs or regressions are introduced (for example, the instance merging versus max tracking changes here are touching on very different things that should probably be implemented and tested separately).

If you prefer to have a chat in realtime, just drop me an email --> talmo@salk.edu.

Thanks again for all the work! Looking forward to getting this merged in :)

Talmo

getzze commented 2 months ago

Now I understand the use of Local queues.

Maybe the best approach is to make 2 PR, one for bringing the max_tracks feature to the Window_size candidate maker and another one to add the merging option. But the name flowmaxtracks would need to be changed in my opinion, to better reflect that it's using a local queue.

I will try to send a screenshot to explain what I mean by a track becoming empty. But basically, when a track is not detected for window_size timesteps, no instance is never assigned to it anymore, ever. I think it has to do with the loop in FlowMaxTracksCandidateMaker.get_candidates, line 408. But I didn't investigate much. When using pre-culling to 2 instances, it would not happen.

I didn't realize that you could limit the number of instances from inference, thanks. Although I like to try different tracking method with new videos, so I run inference and tracking separately and I would still keep all the instances in the inference stage and pre-cull when tracking. Maybe other people would do that also, I don't know.

getzze commented 2 months ago

Regarding the merging function:

For example: if the two node sets are exactly disjoint, but correctly belong to two different animals, we'd be incorrectly merging them into what we call a "chimera" (top of one mouse connected to the bottom of another mouse).

Yes but this chimera will be less likely to be the best match for the track (but it's not sure).

Another scenario we see often that this wouldn't quite resolve is when we have two almost disjoint sets, except for one node that is maybe off by a couple of pixels. Adding a bit of fuzziness to the keypoint matching to test for "disjoint-ness" before triggering this function could help cover those cases as well.

I also see that. I only implemented merging totally disjoint nodes, but we could think of merging group of nodes with one overlapping node (taking the average position of the duplicate node). It would not be too hard to add.

The thing I'm a bit more confused by is that you're mentioning this as useful when there are two instances on the same frame that also belong to the same track. How is this even possible given any of the current tracking approaches? We always do matching that ensures that no more than one instance is assigned to a given track within a frame.

The problem I tried to solve with this merging function is the following:

  1. the 2 mice are close-by, 2 full instances correctly found, say male and female.
  2. next timestep, one instance (male) partially or completely disappears
  3. next timestep, the female instance is split in two: front nodes and back nodes. Because only one partial instance can be assigned to the female track, say the front nodes are assigned to female. the back nodes are not assigned to any track. The male instance is still partially or completely absent.
  4. after some timesteps like this, the back nodes of the "female" is assigned to the male track. The front is still female.
  5. then the front disappears, no instance is assigned to female. The back of the "female" is still male.
  6. the full "female" is assigned to male.

So the idea of merging instances was to avoid the step 3, stopping the identity swap.