The model is based on https://coix.readthedocs.io/en/latest/examples/bmnist.html
The simulator is taken from https://github.com/kirkegaardlab/deeptangle
Data generation is in gen_worm_data.ipynb
The model, inference program and training code are in wormsim_encoder_only_2.ipynb
The model samples
- Worm coordinates from the simulator (6 (x,y)-coordinates per worm)
- Locations in the image
- Generates 28 x 28 frames for each worm for all time steps (interpolating the coordinates from 1)
- Places the frames at the locations sampled in 2.
The inference is amortized population Gibbs sampling (http://proceedings.mlr.press/v119/wu20h/wu20h.pdf), with
- A kernel per time step that proposes, z_where_d_t, the location of each worm d at time t - all kernels are parametrized by the same encoder_where NN, which learns to map from the 64 x 64 frames to the mean and variance of the latent locations
- A kernel per time step that proposes z_what_d_t, a latent variable for the shape of worm d at time t, which is interpreted as coordinates and rendered by the decoder
- A kernel that proposes the simulator parameters, parametrized by a GRU, which learns to invert the simulator, i.e. map from the z_what's to the simulator parameters that produce that worm
Model is still far from convergence, longer training run is in progress:
On the left is the original video, with frames over the inferred locations.
On the right is reconstructed video.
While detection and tracking is fairly reliable, reconstruction is not quite there yet:
