DOCSEARCHbyAlgolia
Submit Docs

TrackNetV3

repository·master·Indexed 43 days ago

https://github.com/qaz812345/tracknetv3

A specialized shuttlecock tracking system utilizing a two-module architecture consisting of a trajectory prediction module and a rectification module. It employs background estimation and inpainting to manage occlusions and visual interference. Built for Ubuntu 16.04.7 LTS, Python 3.8.7, and torch 1.10.0.

Tokens
1.5K
Snippets
7
Records
7
Agent score
6%

What's inside TrackNetV3

  1. Train the Tracking Module

    master

    The tracking module can be trained from scratch or resumed from a previous checkpoint.

    Train from scratch

    Use train.py with the TrackNet model name. Key arguments include --seq_len, --epochs, --batch_size, and --bg_mode (e.g., concat).

    python train.py --model_name TrackNet --seq_len 8 --epochs 30 --batch_size 10 --bg_mode concat --alpha 0.5 --save_dir exp --verbose

    Resume training

    To continue from the last epoch to a specific epoch, use the --resume_training flag.

    python train.py --model_name TrackNet --epochs 30 --save_dir exp --resume_training --verbose

  2. Install TrackNetV3

    master

    To set up the development environment for TrackNetV3, ensure you are using Ubuntu 16.04.7 LTS, Python 3.8.7, and torch 1.10.0.

    Follow these steps to install:

    1. Clone the repository.
    2. Install the required dependencies using pip.
    git clone https://github.com/qaz812345/TrackNetV3.git
cd TrackNetV3
pip install -r requirements.txt

  3. Use the Error Analysis Interface

    master

    TrackNetV3 includes a Dash-based web interface for detailed error analysis.

    1. Generate Detailed Results

    Run evaluation with the --output_pred flag to save detailed JSON results:

    python test.py --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir eval --output_pred

    2. Configure the File List

    Open error_analysis.py and add the JSON paths of your evaluation results to the appropriate eval_file_list (for train, val, or test splits) following this format:

    {'label': label_name, 'value': json_path}

    3. Launch the Dash App

    Run the following command to start the interface:

    python error_analysis.py --split test --host 127.0.0.1

  4. Prepare Dataset for Training

    master

    Before training, you must organize the Shuttlecock Trajectory Dataset into a specific structure.

    Directory Structure

    Your data directory should look like this:

      data
    ├─ train
    |   ├── match1/
    |   |   ├── csv/
    |   |   ├── frame/
    |   |   └── video/
    |   └── ...
    ├─ val
    |   └── ...
    └─ test
        └── ...

    Steps

    1. Merge and Rename: Merge Professional and Amateur match directories into train. Rename Amateur matches to start from match24 through match26. Rename Test to test.
    2. CSV Format: Ensure CSV files contain Frame, Visibility, X, Y.
    3. Preprocessing: Run the preprocessing script to generate frame directories and the val directory.
    4. Configure Data Root: In dataset.py, set the data root directory to your data_dir.

    Note: If you modify dataset.py, delete the generated .npy cache files to ensure changes take effect.

    Background Estimation: The system uses median images from <data_dir>/median. If a match-level median is unavailable, it falls back to the rally-level median.

    python preprocess.py

  5. Train the Rectification Module

    master

    The rectification module (InpaintNet) requires predicted trajectories and inpainting masks generated by the tracking module.

    1. Generate Mask Data

    First, generate the necessary training data for the rectification module:

    python generate_mask_data.py --tracknet_file ckpts/TrackNet_best.pt --batch_size 16

    2. Train InpaintNet

    Train the module from scratch using train.py with the InpaintNet model name.

    python train.py --model_name InpaintNet --seq_len 16 --epoch 300 --batch_size 32 --lr_scheduler StepLR --mask_ratio 0.3 --save_dir exp --verbose

    To resume training, use the --resume_training flag.

  6. Evaluate TrackNetV3

    master

    Evaluation can be performed on the full model (Tracking + Rectification) or just the tracking module.

    Full Model Evaluation

    1. Generate mask data for the test split: python generate_mask_data.py --tracknet_file ckpts/TrackNet_best.pt --split_list test
    2. Run the test script: python test.py --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir eval

    Tracking Module Only

    python test.py --tracknet_file ckpts/TrackNet_best.pt --save_dir eval

    Visualizing Results

    To generate a video comparing ground truth labels with predicted results: python test.py --tracknet_file ckpts/TrackNet_best.pt --video_file data/test/match1/video/1_05_02.mp4

    python generate_mask_data.py --tracknet_file ckpts/TrackNet_best.pt --split_list test
python test.py --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir eval

  7. Run Inference with TrackNetV3

    master

    Use predict.py to generate shuttlecock trajectory labels from a video file. You must provide both the tracking module checkpoint (--tracknet_file) and the inpainting module checkpoint (--inpaintnet_file).

    Basic Usage

    Predict labels and save them to a directory:

    python predict.py --video_file test.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction

    To also output a video file showing the predicted trajectory, add the --output_video flag.

    Handling Large Videos

    For videos that cause memory issues, use the following flags:

    • --large_video: Enables IterableDataset to prevent memory errors (note: this reduces inference speed).
    • --max_sample_num: Sets the number of samples used for background estimation.
    • --video_range: Specifies the start and end seconds (e.g., 324,330) for background estimation.
    python predict.py --video_file test.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction --large_video --video_range 324,330