The codes are based on PSI-Trajectory-Prediction and SGNet.

1. GETTING STARTED

1.0. Clone the repository

git clone https://github.com/KaAI-KMU/KaAI-PSI-Trajectory-Prediction.git

1.1. Install dependencies

Create conda environment.

conda create -n {env_name} python=3.8
conda activate {env_name}

Install pytorch. Please refer to pytorch for the details of installation.

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Install other dependencies.

cd KaAI-PSI-Trajectory-Prediction
pip install -r requirements.txt

1.2. Download data

Please refer to PSI dataset for the details of PSI dataset and data structure.

You can download the pre-processed and structured data(including center optical flow) from here and extract it into KaAI-PSI-Trajectory-Prediction/psi_dataset.

The optical flows are generated by using FlowFormer model which is pretrained with Sintel. And we use optical flow of the center of the bounding box.

You can download the center optical flow data from here and extract it into KaAI-PSI-Trajectory-Prediction/psi_dataset.

Or you can generate the optical flow by yourself. Please refer to FlowFormer and locate the optical flow data into KaAI-PSI-Trajectory-Prediction/psi_dataset/optical_flow.

You can generate the center optical flow by using generate_center_of.py.

python generate_center_of.py

The data structure should be as follows:

- KaAI-PSI-Trajectory
    - psi_dataset
        - PSI2.0_TrainVal
            - annotations
            - splits
        - center_of
            - video_0001
            - video_0002
            - ...

For reducing the time for loading the data, you can download the preprocessed data from here and extract it into KaAI-PSI-Trajectory-Prediction/database. If you want to process the data by yourself, just skip this step.

However, you should download traj_database_test.pkl because it won't be generated by running main.py.

The data structure should be as follows:

- KaAI-PSI-Trajectory
    - database
        - traj_database_train.pkl
        - traj_database_val.pkl
        - traj_database_trainval.pkl
        - traj_database_test.pkl

1.3. Download pre-trained models

We used the pre-trained model of SGNet which is trained with JAAD dataset.

You can download the pre-trained models from here and extract it into root directory(KaAI-PSI-Trajectory).

2. TRAINING

2.1. Train the model

Train the model with training set and validate with validation set.(for dry-run)

python main.py --config_file cfgs/PSI2.0/sgnet_cvae_flow.yaml --train --pretrained_path SGNet_pretrained_with_JAAD.pth --extra_tag {extra_tag} --epochs 1

The model will converge after 1 epoch if you use the pre-trained model.

The results will be saved in KaAI-PSI-Trajectory-Prediction/ckpts/ped_traj/PSI2.0/SGNetCVAETrajBbox/{extra_tag}.

Train the model with full dataset.(for test)

python main.py --config_file cfgs/PSI2.0/sgnet_cvae_flow.yaml --trainval --pretrained_path SGNet_pretrained_with_JAAD.pth --extra_tag {extra_tag} --epochs 1

You can see more options in opts.py

3. EVALUATION

3.1. Evaluate the model with validation set

As default, the model will be evaluated when the training is finished.

However, if you want to evaluate the model with the specific checkpoint with validation set, use --pretrained_path option.

python main.py --config_file cfgs/PSI2.0/sgnet_cvae_flow.yaml --val --pretrained_path {pretrained_path} --extra_tag {extra_tag}

The results will be saved in KaAI-PSI-Trajectory-Prediction/ckpts/ped_traj/PSI2.0/SGNetCVAETrajBbox/{extra_tag}.

3.2. Test the model

If you want to test the model with test set, use --test option.

python main.py --config_file cfgs/PSI2.0/sgnet_cvae_flow.yaml --test --pretrained_path {pretrained_path} --extra_tag {extra_tag}

The results will be saved in KaAI-PSI-Trajectory-Prediction/ckpts/ped_traj/PSI2.0/SGNetCVAETrajBbox/{extra_tag}.