Marker Tracking with Neural Networks

This project provides robust marker tracking for camera-based tactile sensors with Sim2Real neural networks.

Getting Started

System: Ubuntu, OSX

Install packages

pip install tensorflow-gpu opencv-python

Download models

bash ./models/download_model.sh

Download test GelSight video

bash ./data/download_video.sh

Marker tracking for GelSight video

python example_tracking_video.py

Note: The model is implemented initially with Keras/TensorFlow to be compatible with Coral, using TPU as USB Accelerator for Raspberry Pi on-device computation. Please feel free to switch the model to other frameworks, like PyTorch, for your purpose.

Examples

`python example_tracking_sim.py`

The interactive tracking demo. The mouse distorts (translates and rotates) the marker flow. The yellow arrow shows the marker tracking predictions from the neural network. The model can robustly track markers, even with extreme and complex interactions. The model is trained with 10x14 markers.

The model is also robust to marker sizes and background disturbances, due to added domain randomization during training.

User interface:

Mouse Click and Drag": Translate the flow.
Key s and Drag mouse: Rotate the flow. (The mouse needs to be pressed while hitting s.)
Key c: Random background.
Key z: Change marker sizes.
Key r: Reset.

`python example_tracking_sim_generic.py`

The generic model is trained on variable grid patterns so that it can be invariant to different numbers of markers. The output is the flow with the same size of the input.

Additional user interface:

Key d: Random marker patterns.

Note: We suggest trying the generic model for preliminary experiments, and training your fixed model for best performance. The generic one can work on more cases directly, and the fixed one is more accurate for a certain marker pattern.

`python example_tracking_video.py`

The model can be transferred to real sensor data robustly, with large forces, multiple contacts, and wrinkles.

Train

You can train your own model to optimize the performance for a specific marker pattern. Here are two training example that is trained on 10x14 markers (train.py), and variable marker patterns (train_generic.py). Please feel free to customize them for your purposes.

`python train.py`

It takes data pairs (X, Y) from generate_img() in generate_data.py.

Input: X is the synthesized image. It has random backgrounds, and 10x14 markers. The marker positions are randomly distorted (translated and rotated) with Gaussian-based smoothing. It imitates smooth elastomer distortions with simple operations. The markers are then rendered as regions darker than surrounding pixels, given the distorted positions.
Output: Y has a dimension of 10x14x2, which represents the ground-truth horizontal and vertical displacement in pixel for each marker.
Model: The model consists multiple Convolutional layers and Pooling layers, defined as build_model_small() in train.py.

`python train_generic.py`

It takes the data pairs (X, Y) from generate_img() in generate_data_generic.py.

Input: X is generated similarly to the previous one. The difference is that the number of rows and columns is randomly selected from [4, 15].
Output: Y is the flow with the same width and height as the input image. Each point represents the horizontal and vertical displacement for each pixel. We use multiple resolutions of Y to accelerate training.
Model: The model is in the Encoder-Decoder style, which consists of Convolutional, Pooling, and Upsampling layers to generate output with the same dimensions as the input. It is defined as build_model_ae in train.py

Note: Please customize the generate_img() function to fit the marker patterns to your sensor.

Citation

If you use this code for your research, please cite our paper: Gelsight Wedge: Measuring High-Resolution 3D Contact Geometry with a Compact Robot Finger:

@inproceedings{wang2021gelsight,
  title={Gelsight wedge: Measuring high-resolution 3d contact geometry with a compact robot finger},
  author={Wang, Shaoxiong and She, Yu and Romero, Branden and Adelson, Edward},
  booktitle={2021 IEEE International Conference on Robotics and Automation (ICRA)},
  pages={6468--6475},
  year={2021},
  organization={IEEE}
}

wx405557858 / neural_tracking_data