Pytorch implementation based on TrackNetv2.
Supported logging options:
git clone https://github.com/mareksubocz/TrackNet
cd /TrackNet
pip install -r requirements.txt
python train.py --dataset PATH_TO_DATASET --device cuda
python predict.py PATH_TO_VIDEO --weights PATH_TO_TRAINED_WEIGHTS --device cuda
Keybindings:
- l / → : next frame
- h / ← : previous frame
- v : annotate well-visible ball
- o : annotate occluded ball
- m : annotate ball in motion (blurred)
- f : fast-forward/pause video
- n : go to next annotated frame
- x : remove annotation
- = / + : enlarge the annotation mark size
- - : reduce the annotation mark size
- q : finish annotating and save results
python labellingTool.py video.mp4
Labelling tool in use. Fast-forward function is distorted due to gif compression.
| Argument name | Type | Default value | Description |
|---|---|---|---|
| weights | str | None | Path to initial weights the model should be loaded with. If not specified, the model will be initialized with random weights. |
| checkpoint | str | None | Path to a checkpoint, chekpoint differs from weights by to including information about current loss, epoch and optimizer state. |
| batch_size | int | 2 | Batch size of the training dataset. |
| val_batch_size | int | 1 | Batch size of the validation dataset. |
| shuffle | bool | True | Should the dataset be shuffled before training? |
| epochs | int | 10 | Number of epochs. |
| train_size | float | 0.8 | Training dataset size. |
| lr | float | 0.01 | Learning rate. |
| momentum | float | 0.9 | Momentum. |
| dropout | float | 0.0 | Dropout rate. If equals to 0.0, no dropout is used. |
| dataset | str | 'dataset/' | Path to dataset. |
| device | str | 'cpu' | Device to use (cpu, cuda, mps). |
| type | str | 'auto' | Type of dataset to create (auto, image, video). If auto, the dataset type will be inferred from the dataset directory, defaulting to image. |
| save_period | int | 10 | Save checkpoint every x epochs (disabled if <1). |
| save_weights_only | bool | False | Save only weights, not the whole checkpoint |
| save_path | str | 'weights/' | Path to save checkpoints at. |
| no_shuffle | - | - | Don't shuffle the training dataset. |
| tensorboard | - | - | Use tensorboard to log training progress.') |
| one_output_frame | - | - | Demand only one output frame instead of three.') |
| no_save_output_examples | - | - | Don't save output examples to results folder. |
| grayscale | - | - | Use grayscale images instead of RGB.') |
| single_batch_overfit | - | - | Overfit the model on a single batch.') |
Arguments without type or default value are used without an additional value, e.x.
python train.py --dataset dataset/ --grayscale --one_output_frame