paninski-lab / eks

Ensembling and kalman smoothing for pose estimation
MIT License
15 stars 3 forks source link

Unsupervised pose PCA loss used for eks smoothing? #4

Open EugenioBertolini opened 6 months ago

EugenioBertolini commented 6 months ago

Hello, I'm Eugenio Bertolini, a big fan of your works with Lightining Pose (LP) and EKS. Thank you so much for this incredible tool.

In the LP paper you were mentioning that the unsupervised losses that LP outputs (more precisely I'm interested in the pose PCA) could be use by EKS to improve the ensembling, because besides the DLC style confidence, EKS could use additional info to improve its prediction.

Maybe I've misunderstood and this is only possible for the multiview PCA... Are you considering to implement it soon?

Best regards, Eugenio Bertolini Adaptive Motor Control RIKEN Hakubi Research Team

colehurwitz commented 6 months ago

Hi Eugenio,

Thanks for your kind words and for your interest in the EKS repo! Currently, we support

  1. Single-view EKS
  2. Multi-view EKS (using a multi-view PCA as in the LightningPose paper)
  3. Multi-view EKS with asynchronous cameras (mainly for the International Brain Laboratory paw tracking)
  4. Pupil-tracking EKS (assumes the pupil is well-approximated by a diamond)

We haven't implemented the pose PCA EKS yet, but this is certainly an exciting direction! I can discuss with the other team members to see if we can work on this EKS soon.

Already, we find that single-view EKS improves performance a lot over just using a single model. I recommend trying out single-view EKS in the meantime if you have 5-10 models trained.

Best, Cole

EugenioBertolini commented 6 months ago

Thank you for the reply. Yep, I'm running single view eks. With my 5 models and 900,000 frames for each video it takes 5 hours for each label, so I'll get back to you with my results in some days ✌🏻

colehurwitz commented 6 months ago

Ah! We sped this up a lot recently. @keeminlee has been working on this. Maybe we can get this updated for you quickly.

EugenioBertolini commented 6 months ago

Ah cool, thank you so much, that would be great. I was also planning to work on it on my own, but you guys know the project much better so that would bs great 😃

colehurwitz commented 6 months ago

If you have some ideas for speedups, we would be interested to hear about them as well! We can upload our speedups first maybe.

themattinthehatt commented 6 months ago

@EugenioBertolini thanks for the question! How many keypoints per video do you have by the way?

One thing to keep in mind: we are optimizing the smoothing parameter s separately for each keypoint and each video. This optimization procedure is what takes most of the time, as it requires running EKS on the entire video many times. As @colehurwitz mentioned we are working on several ways of speeding this up.

In the meantime, one thing you could do is look at the printouts from some of your videos and record the optimal s values - if these are relatively stable over videos then you could just set the s parameter separately for each keypoint and not run the optimization. This would require you changing the example script a bit to have one smoothing parameter per keypoint, but should be easy enough. I'll let @keeminlee verify this but I believe the optimization usually requires 10-20 EKS runs, meaning a 10x-20x speedup if you don't optimize s.

keeminlee commented 6 months ago

@EugenioBertolini As @themattinthehatt said, the smoothing parameter optimization is responsible for the vast majority of the runtime due to running EKS 10-20 times to arrive at a near-optimal smoothing parameter. Specifying a smoothing parameter can be done via the '--s' flag when running the script.

Another option if you want to still use the optimization is to edit the optimize_smoothing_params() function in eks/core.py to pass in a subset of the frames, which is the y parameter (something like y[:2000] if you want to just use the first 2000 frames to optimize the smoothing parameter).

Thanks for your feedback and please let us know if there are any other questions/ideas!

EugenioBertolini commented 6 months ago

@colehurwitz I checked the code yesterday before running the algorithm, so I am not sure on how I would approach the problem of optimization. Anyway I would have time to work on it from the 16th of May, so I guess you might be faster in updating the package. Thank you for your hard work.

@themattinthehatt Ah got it. I have 31 keypoints, even though I need extra precision only on 6 of them. Now that I know that the optimization is independent for each keypoint, I can adjust to that.

@keeminlee Yeah maybe I will cut the optimization to a few frames only. Good idea. I will tell you how that goes

Thank you for your valuable suggestions. Eugenio Bertolini Adaptive Motor Control RIKEN Hakubi Research Team

EugenioBertolini commented 5 months ago

Hello everyone,

I am getting back on this (after holidays + scholarship application) with good news. The algorithm to estimate the smoothing parameters is quite robust (at least on one of my fly videos).

I have trained LP 5 times, and then I used these 5 models to generate 5 predictions. I estimated the s parameter at 3 different time locations for my LP predicted keypoint. At each time location (of length 16,000 frames), I estimated s one time using all 16,000 frames, two times using 8,000 frames, 4 times using 4,000 frames.

As you can see in the figure, the in total there would be three estimated s parameters for 16,000, six for 8,000, and twelve for 4,000 frames used for the estimation.

The average values of s (black x in the figure), are also reported in this table.

s_dict_avg: { .... all: [8.5528, 9.9782, 8.3767, 8.5132, 5.9158, 6.4468], .... 16000: [8.5536, 9.9731, 8.3259, 8.5121, 5.9142, 6.4458], .... 8000: [8.5528, 9.9761, 8.3521, 8.5119, 5.9151, 6.4470], .... 4000: [8.5521, 9.9854, 8.4523, 8.5157, 5.9181, 6.4476] }

image