This is a fork of UniTR that can be trained and evaluated on data from the CARLA Simulator generated by SimBEV.

UniTR requires the following libraries:

• Python >= 3.8, <3.9
• PyTorch >= 1.1
• nuscenes-dev-kit = 1.0.5
• numpy = 1.23.3
• numba = 0.48.0
• SharedArray = 3.0.0
• shapely = 1.8.4
• llvmlite
• tensorboardX
• easydict
• pyyaml
• scikit-image
• tqdm
• torchvision
• opencv-python
• pyquaternion
• spconv-cu113
• av2
• kornia
• torch_scatter

After installing these dependencies, run

python setup.py develop

to install the codebase.

1. Install Docker on your system.
2. Install the Nvidia Container Toolkit. It exposes your Nvidia graphics card to Docker containers.
3. Install the Nvidia Container Runtime and set it as the default runtime.
4. In the docker folder, run

docker build --no-cache --rm -t unitr:develop .

You may need to replace libnvidia-gl-550 and libnvidia-common-550 packages in the Dockerfile with ones that are compatible with your Nvidia driver version.

The following build arguments (ARG) are available:

• USER: username inside each container, set to untr by default.

Launch a container by running

docker run --privileged --gpus all --network=host -e DISPLAY=$DISPLAY
-v [path/to/UniTR]:/home/unitr
-v [path/to/dataset]:/dataset
--shm-size 32g -it unitr:develop /bin/bash

Then, in /home/unitr, run

python setup.py develop

to install the codebase.

For UniTR 3D object detection, run

bash scripts/dist_train.sh 4 --cfg_file ./cfgs/simbev_models/unitr.yaml --sync_bn --pretrained_model ../unitr_pretrain.pth --logger_iter_interval 50 --num_epochs_to_eval 10 --ckpt_save_time_interval 1800

For UniTR+LSS 3D object detection, run

bash scripts/dist_train.sh 4 --cfg_file ./cfgs/simbev_models/unitr+lss.yaml --sync_bn --pretrained_model ../unitr_pretrain.pth --logger_iter_interval 50 --num_epochs_to_eval 10 --ckpt_save_time_interval 1800

For UniTR BEV segmentation, run

bash scripts/dist_train.sh 4 --cfg_file ./cfgs/simbev_models/unitr_map.yaml --sync_bn --eval_map --logger_iter_interval 50 --num_epochs_to_eval 20 --ckpt_save_time_interval 1800

For UniTR+LSS BEV segmentation, run

bash scripts/dist_train.sh 4 --cfg_file ./cfgs/simbev_models/unitr_map+lss.yaml --sync_bn --eval_map --logger_iter_interval 50 --num_epochs_to_eval 20 --ckpt_save_time_interval 1800

Replace 4 in scripts/dist_train.sh 4 with the number of your GPUs. You can use --batch_size to specify the batch size (default 24) based on your available GPU memory.

If you want to evaluate on the test set instead of the val set, change test under INFO_PATH in cfgs/simbev_models/[model].yaml to ${infos/simbev_infos_test.json}.

For UniTR 3D object detection, run

bash scripts/dist_test.sh 4 --cfg_file ./cfgs/simbev_models/unitr.yaml --ckpt [checkpoint file]

For UniTR+LSS 3D object detection, run

bash scripts/dist_test.sh 4 --cfg_file ./cfgs/simbev_models/unitr+lss.yaml --ckpt [checkpoint file]

For UniTR BEV Segmentation, run

bash scripts/dist_test.sh 4 --cfg_file ./cfgs/simbev_models/unitr_map.yaml --ckpt [checkpoint file] --eval_map

For UniTR+LSS BEV Segmentation, run

bash scripts/dist_test.sh 4 --cfg_file ./cfgs/simbev_models/unitr_map+lss.yaml --ckpt [checkpoint file] --eval_map

SDS: 61.7% / Checkpoint

SDS: 62.2% / Checkpoint

Results are provided for different IoU thresholds.

Checkpoint

Checkpoint