The code in this repository is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.
GolfDB is a high-quality video dataset created for general recognition applications in the sport of golf, and specifically for the task of golf swing sequencing.
This repo contains a simple PyTorch implemention of the SwingNet baseline model presented in the paper. The model was trained on split 1 without any data augmentation and achieved an average PCE of 71.5% (PCE of 76.1% reported in the paper is credited to data augmentation including horizontal flipping and affine transformations).
If you use this repo please cite the GolfDB paper:
@InProceedings{McNally_2019_CVPR_Workshops,
author = {McNally, William and Vats, Kanav and Pinto, Tyler and Dulhanty, Chris and McPhee, John and Wong, Alexander},
title = {GolfDB: A Video Database for Golf Swing Sequencing},
booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
month = {June},
year = {2019}
}Run generate_splits.py to convert the .mat dataset file to a dataframe and generate the 4 splits.
I have provided the preprocessed video clips for a frame size of 160x160 (download here). Place 'videos_160' in the data directory. If you wish to use a different input configuration you must download the YouTube videos (URLs provided in dataset) and preprocess the videos yourself. I have provided preprocess_videos.py to help with that.
Download the MobileNetV2 pretrained weights from this repository and place 'mobilenet_v2.pth.tar' in the root directory.
Run train.py
Train your own model by following the steps above or download the pre-trained weights here. Create a 'models' directory if not already created and place 'swingnet_1800.pth.tar' in this directory.
Run eval.py. If using the pre-trained weights provided, the PCE should be 0.715.
Follow steps above to download pre-trained weights. Then in the terminal: python3 test_video.py -p test_video.mp4
Note: This code requires the sample video to be cropped and cut to bound a single golf swing. I used online video cropping and cutting tools for my golf swing video. See test_video.mp4 for reference.
Good luck!