├── .gitignore ├── LICENSE ├── README.md ├── assets ├── MOT-BDD.png ├── MOT.png ├── MOTS-BDD.png ├── MOTS.png ├── SOT.png ├── Unicorn.png ├── VOS.png ├── data.md ├── install.md ├── model_zoo.md ├── test.md ├── train.md └── video_demo.gif ├── datasets └── data_path │ ├── citypersons.train │ └── eth.train ├── exps └── default │ ├── unicorn_det_convnext_large_800x1280.py │ ├── unicorn_det_convnext_tiny_800x1280.py │ ├── unicorn_det_r50_800x1280.py │ ├── unicorn_inst_convnext_tiny_800x1280.py │ ├── unicorn_track_large.py │ ├── unicorn_track_large_mask.py │ ├── unicorn_track_large_mot_challenge.py │ ├── unicorn_track_large_mot_challenge_mask.py │ ├── unicorn_track_r50.py │ ├── unicorn_track_r50_mask.py │ ├── unicorn_track_tiny.py │ ├── unicorn_track_tiny_mask.py │ ├── unicorn_track_tiny_mot_only.py │ ├── unicorn_track_tiny_mots_only.py │ ├── unicorn_track_tiny_rt.py │ ├── unicorn_track_tiny_rt_mask.py │ ├── unicorn_track_tiny_sot_only.py │ └── unicorn_track_tiny_vos_only.py ├── external ├── PyDavis16EvalToolbox │ ├── .gitignore │ ├── LICENSE │ ├── eval.py │ ├── metrics │ │ ├── __init__.py │ │ ├── f_boundary.py │ │ └── jaccard.py │ └── readme.md ├── davis2017-evaluation │ ├── .gitignore │ ├── LICENSE │ ├── README.md │ ├── davis2017 │ │ ├── __init__.py │ │ ├── davis.py │ │ ├── evaluation.py │ │ ├── metrics.py │ │ ├── results.py │ │ └── utils.py │ ├── evaluation_codalab.py │ ├── evaluation_method.py │ ├── pytest │ │ └── test_evaluation.py │ ├── setup.cfg │ └── setup.py ├── lib │ ├── __init__.py │ ├── test │ │ ├── analysis │ │ │ ├── __init__.py │ │ │ ├── extract_results.py │ │ │ └── plot_results.py │ │ ├── evaluation │ │ │ ├── __init__.py │ │ │ ├── data.py │ │ │ ├── datasets.py │ │ │ ├── environment.py │ │ │ ├── got10kdataset.py │ │ │ ├── lasotdataset.py │ │ │ ├── local.py │ │ │ ├── nfsdataset.py │ │ │ ├── otbdataset.py │ │ │ ├── running.py │ │ │ ├── tc128cedataset.py │ │ │ ├── tc128dataset.py │ │ │ ├── tracker.py │ │ │ ├── trackingnetdataset.py │ │ │ └── uavdataset.py │ │ ├── parameter │ │ │ ├── unicorn_sot.py │ │ │ └── unicorn_vos.py │ │ ├── tracker │ │ │ ├── __init__.py │ │ │ ├── basetracker.py │ │ │ ├── unicorn_sot.py │ │ │ └── unicorn_vos.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── _init_paths.py │ │ │ ├── load_text.py │ │ │ ├── params.py │ │ │ ├── transform_got10k.py │ │ │ └── transform_trackingnet.py │ ├── train │ │ ├── __init__.py │ │ ├── _init_paths.py │ │ ├── admin │ │ │ ├── __init__.py │ │ │ ├── environment.py │ │ │ ├── local.py │ │ │ ├── settings.py │ │ │ └── stats.py │ │ ├── data │ │ │ ├── __init__.py │ │ │ ├── bounding_box_utils.py │ │ │ ├── image_loader.py │ │ │ ├── loader.py │ │ │ ├── processing.py │ │ │ ├── processing_utils.py │ │ │ ├── sampler.py │ │ │ └── transforms.py │ │ └── dataset │ │ │ ├── __init__.py │ │ │ ├── base_video_dataset.py │ │ │ ├── davis.py │ │ │ ├── vos_base.py │ │ │ └── youtubevos.py │ └── utils │ │ ├── __init__.py │ │ ├── box_ops.py │ │ ├── get_image_anno_map.py │ │ ├── helpers.py │ │ ├── lmdb_utils.py │ │ ├── merge.py │ │ ├── misc.py │ │ └── tensor.py └── qdtrack │ ├── .gitignore │ ├── .pre-commit-config.yaml │ ├── LICENSE │ ├── README.md │ ├── clean_seg_track_json.py │ ├── configs │ ├── _base_ │ │ ├── faster_rcnn_r50_fpn.py │ │ └── qdtrack_faster_rcnn_r50_fpn.py │ ├── bdd100k │ │ ├── qdtrack-frcnn_r50_fpn_12e_bdd100k.py │ │ └── unicorn.py │ ├── bdd100k_mots │ │ ├── qdtrack-frcnn_r50_fpn_12e_bdd100k.py │ │ ├── qdtrack-frcnn_r50_fpn_12e_bdd100k_evalseg.py │ │ ├── segtrack-frcnn_r50_fpn_12e_bdd10k.py │ │ ├── segtrack-frcnn_r50_fpn_12e_bdd10k_fixed.py │ │ └── segtrack-frcnn_r50_fpn_12e_bdd10k_fixed_pcan.py │ ├── mot17 │ │ ├── README.md │ │ └── qdtrack-frcnn_r50_fpn_4e_mot17.py │ └── tao │ │ ├── README.md │ │ ├── qdtrack_frcnn_r101_fpn_12e_tao_ft.py │ │ └── qdtrack_frcnn_r101_fpn_24e_lvis.py │ ├── docs │ ├── GET_STARTED.md │ └── INSTALL.md │ ├── download_bdd100k.py │ ├── eval_bdd_submit.sh │ ├── figures │ └── teaser.png │ ├── prepare_bdd100k.sh │ ├── qdtrack │ ├── VERSION │ ├── __init__.py │ ├── apis │ │ ├── __init__.py │ │ ├── inference.py │ │ ├── test.py │ │ ├── test_omni.py │ │ └── train.py │ ├── core │ │ ├── __init__.py │ │ ├── evaluation │ │ │ ├── __init__.py │ │ │ ├── eval_hooks.py │ │ │ ├── mot.py │ │ │ ├── mot_pcan.py │ │ │ └── mots.py │ │ ├── to_bdd100k │ │ │ ├── __init__.py │ │ │ ├── transforms.py │ │ │ └── utils.py │ │ ├── track │ │ │ ├── __init__.py │ │ │ ├── similarity.py │ │ │ ├── transforms.py │ │ │ └── transforms_mots.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ └── visualization.py │ ├── datasets │ │ ├── __init__.py │ │ ├── bdd_video_dataset.py │ │ ├── builder.py │ │ ├── coco_video_dataset.py │ │ ├── mot17_dataset.py │ │ ├── parsers │ │ │ ├── __init__.py │ │ │ ├── coco_api.py │ │ │ └── coco_video_parser.py │ │ ├── pipelines │ │ │ ├── __init__.py │ │ │ ├── formatting.py │ │ │ ├── loading.py │ │ │ └── transforms.py │ │ ├── samplers │ │ │ ├── __init__.py │ │ │ └── distributed_video_sampler.py │ │ └── tao_dataset.py │ ├── models │ │ ├── __init__.py │ │ ├── builder.py │ │ ├── losses │ │ │ ├── __init__.py │ │ │ ├── l2_loss.py │ │ │ └── multipos_cross_entropy_loss.py │ │ ├── mot │ │ │ ├── __init__.py │ │ │ └── qdtrack.py │ │ ├── roi_heads │ │ │ ├── __init__.py │ │ │ ├── quasi_dense_roi_head.py │ │ │ └── track_heads │ │ │ │ ├── __init__.py │ │ │ │ └── quasi_dense_embed_head.py │ │ └── trackers │ │ │ ├── __init__.py │ │ │ ├── quasi_dense_embed_tracker.py │ │ │ └── tao_tracker.py │ ├── utils │ │ ├── __init__.py │ │ ├── collect_env.py │ │ └── logger.py │ └── version.py │ ├── remap_cocofmt.py │ ├── requirements.txt │ ├── setup.cfg │ ├── setup.py │ └── tools │ ├── convert_datasets │ ├── mot2coco.py │ └── tao2coco.py │ ├── debug_data.py │ ├── debug_test.py │ ├── dist_debug_test.sh │ ├── dist_test.sh │ ├── dist_test_omni.sh │ ├── dist_train.sh │ ├── eval.py │ ├── slurm_test.sh │ ├── slurm_train.sh │ ├── test.py │ ├── test_omni.py │ ├── to_bdd100k.py │ └── train.py ├── launch_uni.py ├── requirements.txt ├── setup.cfg ├── setup.py ├── tools ├── __init__.py ├── _init_paths.py ├── analysis_results.py ├── convert_cityperson_to_coco.py ├── convert_crowdhuman_to_coco.py ├── convert_ethz_to_coco.py ├── convert_mot17_to_coco.py ├── convert_mot17_to_omni.py ├── convert_mot20_to_coco.py ├── convert_mots_to_coco.py ├── copy_1to3.py ├── demo.py ├── eval.py ├── export_torchscript.py ├── interpolation.py ├── process_trackingnet.py ├── test.py ├── track.py ├── track_omni.py ├── train.py └── train_dist.py └── unicorn ├── __init__.py ├── core ├── __init__.py ├── launch.py └── trainer.py ├── data ├── __init__.py ├── data_augment.py ├── data_prefetcher.py ├── dataloading.py ├── datasets │ ├── __init__.py │ ├── bdd.py │ ├── bdd_omni.py │ ├── bdd_omni_mots.py │ ├── coco.py │ ├── coco_classes.py │ ├── coco_inst.py │ ├── coco_mots.py │ ├── coco_sot.py │ ├── data_specs │ │ ├── got10k_vot_train_split.txt │ │ ├── lasot_train_split.txt │ │ └── trackingnet_classmap.txt │ ├── datasets_wrapper.py │ ├── davis.py │ ├── got10k.py │ ├── lasot.py │ ├── mosaicdetection.py │ ├── mosaicdetection_uni.py │ ├── mot.py │ ├── mot_omni.py │ ├── mots_mot.py │ ├── omni_data.py │ ├── saliency.py │ ├── tracking_net.py │ ├── voc.py │ ├── voc_classes.py │ └── youtube_vos.py └── samplers.py ├── evaluators ├── __init__.py ├── bdd_evaluator.py ├── coco_evaluator.py ├── coco_inst_evaluator.py ├── evaluation.py ├── mot_evaluator.py ├── voc_eval.py └── voc_evaluator.py ├── exp ├── __init__.py ├── base_exp.py ├── build.py ├── unicorn_det.py ├── unicorn_det_mask.py ├── unicorn_track.py └── unicorn_track_mask.py ├── layers ├── __init__.py ├── csrc │ ├── cocoeval │ │ ├── cocoeval.cpp │ │ └── cocoeval.h │ └── vision.cpp └── fast_coco_eval_api.py ├── models ├── __init__.py ├── backbone │ ├── convnext.py │ ├── darknet.py │ ├── network_blocks.py │ ├── resnet.py │ ├── swin_config.py │ ├── swin_transformer.py │ └── yolo_pafpn_new.py ├── condinst │ ├── __init__.py │ ├── comm.py │ ├── config │ │ ├── __init__.py │ │ ├── config.py │ │ ├── defaults.py │ │ └── defaults_d2.py │ ├── conv_with_kaiming_uniform.py │ ├── dynamic_mask_head.py │ └── mask_branch.py ├── deformable_transformer.py ├── losses.py ├── ops │ ├── functions │ │ ├── __init__.py │ │ └── ms_deform_attn_func.py │ ├── make.sh │ ├── modules │ │ ├── __init__.py │ │ └── ms_deform_attn.py │ ├── setup.py │ ├── src │ │ ├── cpu │ │ │ ├── ms_deform_attn_cpu.cpp │ │ │ └── ms_deform_attn_cpu.h │ │ ├── cuda │ │ │ ├── ms_deform_attn_cuda.cu │ │ │ ├── ms_deform_attn_cuda.h │ │ │ └── ms_deform_im2col_cuda.cuh │ │ ├── ms_deform_attn.h │ │ └── vision.cpp │ └── test.py ├── position_encoding.py ├── transformer_encoder.py ├── unicorn.py ├── unicorn_head.py ├── unicorn_head_mask.py ├── yolo_head_det.py ├── yolo_head_det_mask.py └── yolox.py ├── tracker ├── basetrack.py ├── byte_tracker.py ├── kalman_filter.py ├── matching.py └── quasi_dense_embed_tracker.py └── utils ├── __init__.py ├── allreduce_norm.py ├── boxes.py ├── checkpoint.py ├── demo_utils.py ├── dist.py ├── ema.py ├── logger.py ├── lr_scheduler.py ├── merge.py ├── metric.py ├── model_utils.py ├── setup_env.py └── visualize.py /.gitignore: -------------------------------------------------------------------------------- 1 | ### Linux ### 2 | *~ 3 | 4 | # debug, detection and tracking results 5 | debug/ 6 | det/ 7 | test/result_plots 8 | test/tracking_results 9 | test/segmentation_results 10 | video* 11 | Unicorn_outputs/ 12 | 13 | # pyc 14 | *.pyc 15 | 16 | # vscode 17 | *.code-workspace 18 | .vscode 19 | 20 | # vim 21 | .vim 22 | 23 | # dataset 24 | datasets/ 25 | 26 | # checkpoints 27 | *.pth 28 | 29 | # wheels 30 | *.whl 31 | 32 | # images 33 | *.jpg 34 | *.png 35 | 36 | # txts 37 | *.txt 38 | 39 | # C extensions 40 | *.so 41 | 42 | # track results 43 | *.pkl 44 | 45 | # Distribution / packaging 46 | .Python 47 | build/ 48 | develop-eggs/ 49 | dist/ 50 | downloads/ 51 | eggs/ 52 | .eggs/ 53 | # lib/ 54 | lib64/ 55 | parts/ 56 | sdist/ 57 | var/ 58 | wheels/ 59 | pip-wheel-metadata/ 60 | share/python-wheels/ 61 | *.egg-info/ 62 | .installed.cfg 63 | *.egg 64 | MANIFEST 65 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Bin Yan 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /assets/MOT-BDD.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/MOT-BDD.png -------------------------------------------------------------------------------- /assets/MOT.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/MOT.png -------------------------------------------------------------------------------- /assets/MOTS-BDD.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/MOTS-BDD.png -------------------------------------------------------------------------------- /assets/MOTS.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/MOTS.png -------------------------------------------------------------------------------- /assets/SOT.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/SOT.png -------------------------------------------------------------------------------- /assets/Unicorn.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/Unicorn.png -------------------------------------------------------------------------------- /assets/VOS.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/VOS.png -------------------------------------------------------------------------------- /assets/install.md: -------------------------------------------------------------------------------- 1 | # Install 2 | ## Requirements 3 | We test the codes in the following environments, other versions may also be compatible but Pytorch vision should be >= 1.7 4 | 5 | - CUDA 11.3 6 | - Python 3.7 7 | - Pytorch 1.10.0 8 | - Torchvison 0.11.1 9 | 10 | ## Install environment for Unicorn 11 | 12 | ``` 13 | # Pytorch and Torchvision 14 | conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge 15 | 16 | # YOLOX and some other packages 17 | pip3 install -U pip && pip3 install -r requirements.txt 18 | pip3 install -v -e . # or python3 setup.py develop 19 | 20 | # Install Deformable Attention 21 | cd unicorn/models/ops 22 | bash make.sh 23 | cd ../../.. 24 | 25 | # Install mmcv, mmdet, bdd100k 26 | cd external/qdtrack 27 | wget -c https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/mmcv_full-1.4.6-cp37-cp37m-manylinux1_x86_64.whl # This should change according to cuda version and pytorch version 28 | pip3 install --user mmcv_full-1.4.6-cp37-cp37m-manylinux1_x86_64.whl 29 | pip3 install --user mmdet 30 | git clone https://github.com/bdd100k/bdd100k.git 31 | cd bdd100k 32 | python3 setup.py develop --user 33 | pip3 uninstall -y scalabel 34 | pip3 install --user git+https://github.com/scalabel/scalabel.git 35 | cd ../../.. 36 | ``` 37 | -------------------------------------------------------------------------------- /assets/train.md: -------------------------------------------------------------------------------- 1 | # Tutorial for Training 2 | 3 | Every experiment is defined using a python file under the `exps/default` folder. Experiments about object detection and object tracking start with `unicorn_det` and `unicorn_track` respectively. In the next paragraphs, {exp_name} should be replaced with specifc filenames (without .py). For example, if you want to train unicorn with convnext-tiny backbone for object tracking, replace ${exp_name} with unicorn_track_tiny 4 | 5 | ## Detection & Instance Segmentation 6 | 7 | **Single-node Training** 8 | 9 | On a single node with 8 GPUs, run 10 | ``` 11 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 64 --mode multiple --fp16 0 12 | ``` 13 | 14 | **Multiple-node Training** 15 | 16 | On the master node, run 17 | ``` 18 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 128 --mode distribute --fp16 0 --nnodes 2 --master_address ${master_address} --node_rank 0 19 | ``` 20 | 21 | On the second node, run 22 | ``` 23 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 128 --mode distribute --fp16 0 --nnodes 2 --master_address ${master_address} --node_rank 1 24 | ``` 25 | 26 | Testing (Instance Segmentation) 27 | ``` 28 | python3 tools/eval.py -f exps/default/${exp_name}.py -c Unicorn_outputs/${exp_name}/latest_ckpt.pth -b 64 -d 8 --conf 0.001 --mask_thres 0.3 29 | ``` 30 | 31 | ## Unified Tracking (SOT, MOT, VOS, MOTS) 32 | 33 | **Single-node Training** 34 | 35 | On a single node with 8 GPUs, run 36 | ``` 37 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 16 --mode multiple 38 | ``` 39 | **Multiple-node Training** 40 | 41 | On the master node, run 42 | ``` 43 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 32 --mode distribute --nnodes 2 --master_address ${master_address} --node_rank 0 44 | ``` 45 | On the second node, run 46 | ``` 47 | python3 launch_uni.py --name ${exp_name} --nproc_per_node 8 --batch 32 --mode distribute --nnodes 2 --master_address ${master_address} --node_rank 1 48 | ``` 49 | -------------------------------------------------------------------------------- /assets/video_demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/assets/video_demo.gif -------------------------------------------------------------------------------- /exps/default/unicorn_det_convnext_large_800x1280.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpDet 6 | 7 | class Exp(ExpDet): 8 | def __init__(self): 9 | super(Exp, self).__init__() 10 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 11 | self.backbone_name = "convnext_large" 12 | self.pretrained_name = "convnext_large_22k_224.pth" 13 | self.in_channels = [384, 768, 1536] 14 | self.use_checkpoint = True 15 | self.input_size = (800, 1280) 16 | self.test_size = (800, 1280) -------------------------------------------------------------------------------- /exps/default/unicorn_det_convnext_tiny_800x1280.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpDet 6 | 7 | class Exp(ExpDet): 8 | def __init__(self): 9 | super(Exp, self).__init__() 10 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 11 | self.input_size = (800, 1280) 12 | self.test_size = (800, 1280) 13 | -------------------------------------------------------------------------------- /exps/default/unicorn_det_r50_800x1280.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpDet 6 | 7 | class Exp(ExpDet): 8 | def __init__(self): 9 | super(Exp, self).__init__() 10 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 11 | self.backbone_name = "resnet50" 12 | self.in_channels = [512, 1024, 2048] 13 | self.input_size = (800, 1280) 14 | self.test_size = (800, 1280) 15 | -------------------------------------------------------------------------------- /exps/default/unicorn_inst_convnext_tiny_800x1280.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpDetMask 6 | 7 | class Exp(ExpDetMask): 8 | def __init__(self): 9 | super(Exp, self).__init__() 10 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 11 | self.input_size = (800, 1280) 12 | self.test_size = (800, 1280) 13 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_large.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The main setting used in the Unicorn paper (ConvNext-Large Backbone) 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "convnext_large" 15 | self.in_channels = [384, 768, 1536] 16 | self.pretrain_name = "unicorn_det_convnext_large_800x1280" 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_large_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The main setting used in the Unicorn paper (ConvNext-Large Backbone) 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "convnext_large" 15 | self.in_channels = [384, 768, 1536] 16 | self.pretrain_name = "unicorn_track_large" 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_large_mot_challenge.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The main setting used in the Unicorn paper (ConvNext-Large Backbone) 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "convnext_large" 15 | self.in_channels = [384, 768, 1536] 16 | self.pretrain_name = "unicorn_det_convnext_large_800x1280" 17 | self.mot_test_name = "motchallenge" 18 | self.num_classes = 1 19 | self.mhs = False 20 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_large_mot_challenge_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The main setting used in the Unicorn paper (ConvNext-Large Backbone) 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "convnext_large" 15 | self.in_channels = [384, 768, 1536] 16 | self.pretrain_name = "unicorn_track_large_mot_challenge" 17 | self.mot_test_name = "motchallenge" 18 | self.num_classes = 1 19 | self.mhs = False 20 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_r50.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The R50 setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "resnet50" 15 | self.in_channels = [512, 1024, 2048] 16 | self.pretrain_name = "unicorn_det_r50_800x1280" 17 | 18 | 19 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_r50_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The R50 setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.backbone_name = "resnet50" 15 | self.in_channels = [512, 1024, 2048] 16 | self.pretrain_name = "unicorn_track_r50" 17 | 18 | 19 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The baseline setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_det_convnext_tiny_800x1280" 15 | 16 | 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The baseline setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_track_tiny" 15 | 16 | 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_mot_only.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The MOT-only setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_det_convnext_tiny_800x1280" 15 | self.mot_only = True 16 | self.samples_per_epoch = 100000 17 | self.mot_weight = 1 18 | self.mhs = False 19 | 20 | 21 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_mots_only.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The MOT-only setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_track_tiny_mot_only" 15 | self.mot_only = True 16 | self.samples_per_epoch = 100000 17 | self.mot_weight = 1 18 | self.mhs = False 19 | 20 | 21 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_rt.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The Real-Time setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_det_convnext_tiny_800x1280" 15 | self.input_size = (640, 1024) 16 | self.test_size = (640, 1024) 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_rt_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The Real-Time setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_track_tiny_rt" 15 | self.input_size = (640, 1024) 16 | self.test_size = (640, 1024) 17 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_sot_only.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrack 6 | """ 7 | The SOT-only setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrack): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_det_convnext_tiny_800x1280" 15 | self.sot_only = True 16 | self.samples_per_epoch = 100000 17 | 18 | 19 | -------------------------------------------------------------------------------- /exps/default/unicorn_track_tiny_vos_only.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 4 | import os 5 | from unicorn.exp import ExpTrackMask 6 | """ 7 | The VOS-only setting used in the ablation 8 | We load weights pretrained on COCO with input resolution of 800x1280 9 | """ 10 | class Exp(ExpTrackMask): 11 | def __init__(self): 12 | super(Exp, self).__init__() 13 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] 14 | self.pretrain_name = "unicorn_track_tiny_sot_only" 15 | self.sot_only = True 16 | self.samples_per_epoch = 100000 17 | 18 | 19 | -------------------------------------------------------------------------------- /external/PyDavis16EvalToolbox/.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # private 132 | .idea/ 133 | .vscode/ 134 | -------------------------------------------------------------------------------- /external/PyDavis16EvalToolbox/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Pang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /external/PyDavis16EvalToolbox/metrics/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/external/PyDavis16EvalToolbox/metrics/__init__.py -------------------------------------------------------------------------------- /external/PyDavis16EvalToolbox/metrics/jaccard.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # ---------------------------------------------------------------------------- 3 | # A Benchmark Dataset and Evaluation Methodology for Video Object Segmentation 4 | # ----------------------------------------------------------------------------- 5 | # Copyright (c) 2016 Federico Perazzi 6 | # Licensed under the BSD License [see LICENSE for details] 7 | # Written by Federico Perazzi 8 | # ---------------------------------------------------------------------------- 9 | 10 | """ Compute Jaccard Index. """ 11 | 12 | import numpy as np 13 | 14 | 15 | def db_eval_iou(annotation, segmentation): 16 | """Compute region similarity as the Jaccard Index. 17 | 18 | Arguments: 19 | annotation (ndarray): binary annotation map. 20 | segmentation (ndarray): binary segmentation map. 21 | 22 | Return: 23 | jaccard (float): region similarity 24 | 25 | """ 26 | 27 | annotation = annotation.astype(np.bool) 28 | segmentation = segmentation.astype(np.bool) 29 | 30 | if np.isclose(np.sum(annotation), 0) and np.isclose(np.sum(segmentation), 0): 31 | return 1 32 | else: 33 | return np.sum((annotation & segmentation)) / np.sum( 34 | (annotation | segmentation), dtype=np.float32 35 | ) 36 | -------------------------------------------------------------------------------- /external/PyDavis16EvalToolbox/readme.md: -------------------------------------------------------------------------------- 1 | # 基于python的视频目标分割测评工具箱 2 | 3 | A Python-based video object segmentation evaluation toolbox. 4 | 5 | ## 特性 6 | 7 | * 针对**DAVIS 2016无监督视频目标分割**任务，提供`"J(M)", "J(O)", "J(D)", "F(M)", "F(O)", "F(D)"`等指标的评估（代码借鉴自davis官方的代码，建议使用前验证下） 8 | - 导出对指定的模型预测结果的评估结果 9 | - 表格化展示不同视频上模型预测的性能 10 | 11 | ## 依赖项 12 | 13 | - `joblib`: 多进程加速计算 14 | - `prettytable`: 输出好看的表格 15 | 16 | ## 使用方法 17 | 18 | ### DAVIS 2016无监督视频目标分割任务 19 | 20 | 1. `python eval_unvos_method.py --help` 21 | 2. 配置相关项后执行代码 22 | 23 | python3 eval.py --name_list_path /opt/tiger/omnitrack/data/DAVIS/ImageSets/2016/val.txt --mask_root /opt/tiger/omnitrack/data/DAVIS/Annotations/480p --pred_path /opt/tiger/omnitrack/test/vos_baseline_tiny_10f_stm_DAVIS_16val --save_path /opt/tiger/omnitrack/result.pkl -------------------------------------------------------------------------------- /external/davis2017-evaluation/.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | 49 | # Translations 50 | *.mo 51 | *.pot 52 | 53 | # Django stuff: 54 | *.log 55 | local_settings.py 56 | 57 | # Flask stuff: 58 | instance/ 59 | .webassets-cache 60 | 61 | # Scrapy stuff: 62 | .scrapy 63 | 64 | # Sphinx documentation 65 | docs/_build/ 66 | 67 | # PyBuilder 68 | target/ 69 | 70 | # Jupyter Notebook 71 | .ipynb_checkpoints 72 | 73 | # pyenv 74 | .python-version 75 | 76 | # celery beat schedule file 77 | celerybeat-schedule 78 | 79 | # SageMath parsed files 80 | *.sage.py 81 | 82 | # dotenv 83 | .env 84 | 85 | # virtualenv 86 | .venv 87 | venv/ 88 | ENV/ 89 | 90 | # Spyder project settings 91 | .spyderproject 92 | .spyproject 93 | 94 | # Rope project settings 95 | .ropeproject 96 | 97 | # mkdocs documentation 98 | docs/site/ 99 | /site 100 | 101 | # mypy 102 | .mypy_cache/ 103 | 104 | # pytest 105 | .pytest_cache 106 | 107 | # Pylint 108 | .pylintrc 109 | 110 | # PyCharm 111 | .idea/ 112 | .DS_Store 113 | 114 | # Generated C code 115 | _mask.c 116 | -------------------------------------------------------------------------------- /external/davis2017-evaluation/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2020, DAVIS: Densely Annotated VIdeo Segmentation 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | 3. Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /external/davis2017-evaluation/davis2017/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | __version__ = '0.1.0' 4 | -------------------------------------------------------------------------------- /external/davis2017-evaluation/davis2017/results.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from PIL import Image 4 | import sys 5 | 6 | 7 | class Results(object): 8 | def __init__(self, root_dir): 9 | self.root_dir = root_dir 10 | 11 | def _read_mask(self, sequence, frame_id): 12 | try: 13 | mask_path = os.path.join(self.root_dir, sequence, f'{frame_id}.png') 14 | return np.array(Image.open(mask_path)) 15 | except IOError as err: 16 | sys.stdout.write(sequence + " frame %s not found!\n" % frame_id) 17 | sys.stdout.write("The frames have to be indexed PNG files placed inside the corespondent sequence " 18 | "folder.\nThe indexes have to match with the initial frame.\n") 19 | sys.stderr.write("IOError: " + err.strerror + "\n") 20 | sys.exit() 21 | 22 | def read_masks(self, sequence, masks_id): 23 | mask_0 = self._read_mask(sequence, masks_id[0]) 24 | masks = np.zeros((len(masks_id), *mask_0.shape)) 25 | for ii, m in enumerate(masks_id): 26 | masks[ii, ...] = self._read_mask(sequence, m) 27 | num_objects = int(np.max(masks)) 28 | tmp = np.ones((num_objects, *masks.shape)) 29 | tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None] 30 | masks = (tmp == masks[None, ...]) > 0 31 | return masks 32 | -------------------------------------------------------------------------------- /external/davis2017-evaluation/setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = davis2017 3 | version = attr: davis2017.__version__ 4 | description = Evaluation Framework for DAVIS 2017 Semi-supervised and Unsupervised used in the DAVIS Challenges 5 | long_description = file: README.md 6 | long_description_content_type = text/markdown 7 | keywords = segmentation 8 | license = GPL v3 9 | author = Sergi Caelles 10 | author-email = scaelles@vision.ee.ethz.ch 11 | home-page = https://github.com/davisvideochallenge/davis2017-evaluation 12 | classifiers = 13 | Development Status :: 4 - Beta 14 | Intended Audience :: Developers 15 | Intended Audience :: Education 16 | Intended Audience :: Science/Research 17 | License :: OSI Approved :: GNU General Public License v3 (GPLv3) 18 | Programming Language :: Python :: 3.6 19 | Programming Language :: Python :: 3.7 20 | Topic :: Scientific/Engineering :: Human Machine Interfaces 21 | Topic :: Software Development :: Libraries 22 | Topic :: Software Development :: Libraries :: Python Modules 23 | -------------------------------------------------------------------------------- /external/davis2017-evaluation/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | import sys 3 | 4 | if sys.version_info < (3, 6): 5 | sys.exit('Sorry, only Python >= 3.6 is supported') 6 | 7 | setup( 8 | python_requires='>=3.6, <4', 9 | install_requires=[ 10 | 'Pillow>=4.1.1', 11 | 'networkx>=2.0', 12 | 'numpy>=1.12.1', 13 | 'opencv-python>=4.0.0.21', 14 | 'pandas>=0.21.1', 15 | 'pathlib2;python_version<"3.5"', 16 | 'scikit-image>=0.13.1', 17 | 'scikit-learn>=0.18', 18 | 'scipy>=1.0.0', 19 | 'tqdm>=4.28.1' 20 | ]) 21 | -------------------------------------------------------------------------------- /external/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/external/lib/__init__.py -------------------------------------------------------------------------------- /external/lib/test/analysis/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/external/lib/test/analysis/__init__.py -------------------------------------------------------------------------------- /external/lib/test/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from .data import Sequence 2 | from .tracker import Tracker, trackerlist 3 | from .datasets import get_dataset 4 | from .environment import create_default_local_file_test -------------------------------------------------------------------------------- /external/lib/test/evaluation/datasets.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import importlib 3 | from lib.test.evaluation.data import SequenceList 4 | 5 | DatasetInfo = namedtuple('DatasetInfo', ['module', 'class_name', 'kwargs']) 6 | 7 | pt = "lib.test.evaluation.%sdataset" # Useful abbreviations to reduce the clutter 8 | 9 | dataset_dict = dict( 10 | otb=DatasetInfo(module=pt % "otb", class_name="OTBDataset", kwargs=dict()), 11 | nfs=DatasetInfo(module=pt % "nfs", class_name="NFSDataset", kwargs=dict()), 12 | uav=DatasetInfo(module=pt % "uav", class_name="UAVDataset", kwargs=dict()), 13 | tc128=DatasetInfo(module=pt % "tc128", class_name="TC128Dataset", kwargs=dict()), 14 | tc128ce=DatasetInfo(module=pt % "tc128ce", class_name="TC128CEDataset", kwargs=dict()), 15 | trackingnet=DatasetInfo(module=pt % "trackingnet", class_name="TrackingNetDataset", kwargs=dict()), 16 | got10k_test=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='test')), 17 | got10k_val=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='val')), 18 | got10k_ltrval=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='ltrval')), 19 | lasot=DatasetInfo(module=pt % "lasot", class_name="LaSOTDataset", kwargs=dict()), 20 | lasot_lmdb=DatasetInfo(module=pt % "lasot_lmdb", class_name="LaSOTlmdbDataset", kwargs=dict()), 21 | dv2017_val=DatasetInfo(module="lib.train.dataset.davis", class_name="Davis", kwargs=dict(version='2017', split='val')), 22 | dv2016_val=DatasetInfo(module="lib.train.dataset.davis", class_name="Davis", kwargs=dict(version='2016', split='val')), 23 | yt2018_valid=DatasetInfo(module="lib.train.dataset.youtubevos", class_name="YouTubeVOS", 24 | kwargs=dict(version='2018', split='val')), 25 | ) 26 | 27 | 28 | def load_dataset(name: str): 29 | """ Import and load a single dataset.""" 30 | name = name.lower() 31 | dset_info = dataset_dict.get(name) 32 | if dset_info is None: 33 | raise ValueError('Unknown dataset \'%s\'' % name) 34 | 35 | m = importlib.import_module(dset_info.module) 36 | dataset = getattr(m, dset_info.class_name)(**dset_info.kwargs) # Call the constructor 37 | return dataset.get_sequence_list() 38 | 39 | 40 | def get_dataset(*args): 41 | """ Get a single or set of datasets.""" 42 | dset = SequenceList() 43 | for name in args: 44 | dset.extend(load_dataset(name)) 45 | return dset -------------------------------------------------------------------------------- /external/lib/test/evaluation/got10kdataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList 3 | from lib.test.utils.load_text import load_text 4 | import os 5 | 6 | 7 | class GOT10KDataset(BaseDataset): 8 | """ GOT-10k dataset. 9 | 10 | Publication: 11 | GOT-10k: A Large High-Diversity Benchmark for Generic Object Tracking in the Wild 12 | Lianghua Huang, Xin Zhao, and Kaiqi Huang 13 | arXiv:1810.11981, 2018 14 | https://arxiv.org/pdf/1810.11981.pdf 15 | 16 | Download dataset from http://got-10k.aitestunion.com/downloads 17 | """ 18 | def __init__(self, split): 19 | super().__init__() 20 | # Split can be test, val, or ltrval (a validation split consisting of videos from the official train set) 21 | if split == 'test' or split == 'val': 22 | self.base_path = os.path.join(self.env_settings.got10k_path, split) 23 | else: 24 | self.base_path = os.path.join(self.env_settings.got10k_path, 'train') 25 | 26 | self.sequence_list = self._get_sequence_list(split) 27 | self.split = split 28 | 29 | def get_sequence_list(self): 30 | return SequenceList([self._construct_sequence(s) for s in self.sequence_list]) 31 | 32 | def _construct_sequence(self, sequence_name): 33 | anno_path = '{}/{}/groundtruth.txt'.format(self.base_path, sequence_name) 34 | 35 | ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64) 36 | 37 | frames_path = '{}/{}'.format(self.base_path, sequence_name) 38 | frame_list = [frame for frame in os.listdir(frames_path) if frame.endswith(".jpg")] 39 | frame_list.sort(key=lambda f: int(f[:-4])) 40 | frames_list = [os.path.join(frames_path, frame) for frame in frame_list] 41 | 42 | return Sequence(sequence_name, frames_list, 'got10k', ground_truth_rect.reshape(-1, 4)) 43 | 44 | def __len__(self): 45 | return len(self.sequence_list) 46 | 47 | def _get_sequence_list(self, split): 48 | with open('{}/list.txt'.format(self.base_path)) as f: 49 | sequence_list = f.read().splitlines() 50 | 51 | if split == 'ltrval': 52 | with open('{}/got10k_val_split.txt'.format(self.env_settings.dataspec_path)) as f: 53 | seq_ids = f.read().splitlines() 54 | 55 | sequence_list = [sequence_list[int(x)] for x in seq_ids] 56 | return sequence_list 57 | -------------------------------------------------------------------------------- /external/lib/test/evaluation/local.py: -------------------------------------------------------------------------------- 1 | from lib.test.evaluation.environment import EnvSettings 2 | from unicorn.data import get_unicorn_datadir 3 | import os 4 | 5 | def local_env_settings(): 6 | settings = EnvSettings() 7 | # Set your local paths here. 8 | settings.prj_dir = os.path.join(get_unicorn_datadir(), "..") 9 | settings.data_dir = get_unicorn_datadir() 10 | settings.got10k_path = os.path.join(settings.data_dir, 'GOT10K') 11 | settings.lasot_path = os.path.join(settings.data_dir, 'LaSOT') 12 | settings.trackingnet_path = os.path.join(settings.data_dir, 'TrackingNet') 13 | settings.nfs_path = os.path.join(settings.data_dir, 'nfs') 14 | settings.otb_path = os.path.join(settings.data_dir, 'OTB2015') 15 | settings.uav_path = os.path.join(settings.data_dir, 'UAV123') 16 | settings.vot_path = os.path.join(settings.data_dir, 'VOT2019') 17 | settings.result_plot_path = os.path.join(settings.prj_dir, 'test/result_plots') 18 | settings.results_path = os.path.join(settings.prj_dir, 'test/tracking_results') # Where to store tracking results 19 | settings.segmentation_path = os.path.join(settings.prj_dir, 'test/segmentation_results') 20 | 21 | return settings 22 | 23 | -------------------------------------------------------------------------------- /external/lib/test/evaluation/tc128cedataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList 3 | import os 4 | import glob 5 | import six 6 | 7 | 8 | class TC128CEDataset(BaseDataset): 9 | """ 10 | TC-128 Dataset (78 newly added sequences) 11 | modified from the implementation in got10k-toolkit (https://github.com/got-10k/toolkit) 12 | """ 13 | def __init__(self): 14 | super().__init__() 15 | self.base_path = self.env_settings.tc128_path 16 | self.anno_files = sorted(glob.glob( 17 | os.path.join(self.base_path, '*/*_gt.txt'))) 18 | """filter the newly added sequences (_ce)""" 19 | self.anno_files = [s for s in self.anno_files if "_ce" in s] 20 | self.seq_dirs = [os.path.dirname(f) for f in self.anno_files] 21 | self.seq_names = [os.path.basename(d) for d in self.seq_dirs] 22 | # valid frame range for each sequence 23 | self.range_files = [glob.glob(os.path.join(d, '*_frames.txt'))[0] for d in self.seq_dirs] 24 | 25 | def get_sequence_list(self): 26 | return SequenceList([self._construct_sequence(s) for s in self.seq_names]) 27 | 28 | def _construct_sequence(self, sequence_name): 29 | if isinstance(sequence_name, six.string_types): 30 | if not sequence_name in self.seq_names: 31 | raise Exception('Sequence {} not found.'.format(sequence_name)) 32 | index = self.seq_names.index(sequence_name) 33 | # load valid frame range 34 | frames = np.loadtxt(self.range_files[index], dtype=int, delimiter=',') 35 | img_files = [os.path.join(self.seq_dirs[index], 'img/%04d.jpg' % f) for f in range(frames[0], frames[1] + 1)] 36 | 37 | # load annotations 38 | anno = np.loadtxt(self.anno_files[index], delimiter=',') 39 | assert len(img_files) == len(anno) 40 | assert anno.shape[1] == 4 41 | 42 | # return img_files, anno 43 | return Sequence(sequence_name, img_files, 'tc128', anno.reshape(-1, 4)) 44 | 45 | def __len__(self): 46 | return len(self.seq_names) 47 | -------------------------------------------------------------------------------- /external/lib/test/evaluation/tc128dataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList 3 | import os 4 | import glob 5 | import six 6 | 7 | 8 | class TC128Dataset(BaseDataset): 9 | """ 10 | TC-128 Dataset 11 | modified from the implementation in got10k-toolkit (https://github.com/got-10k/toolkit) 12 | """ 13 | def __init__(self): 14 | super().__init__() 15 | self.base_path = self.env_settings.tc128_path 16 | self.anno_files = sorted(glob.glob( 17 | os.path.join(self.base_path, '*/*_gt.txt'))) 18 | self.seq_dirs = [os.path.dirname(f) for f in self.anno_files] 19 | self.seq_names = [os.path.basename(d) for d in self.seq_dirs] 20 | # valid frame range for each sequence 21 | self.range_files = [glob.glob(os.path.join(d, '*_frames.txt'))[0] for d in self.seq_dirs] 22 | 23 | def get_sequence_list(self): 24 | return SequenceList([self._construct_sequence(s) for s in self.seq_names]) 25 | 26 | def _construct_sequence(self, sequence_name): 27 | if isinstance(sequence_name, six.string_types): 28 | if not sequence_name in self.seq_names: 29 | raise Exception('Sequence {} not found.'.format(sequence_name)) 30 | index = self.seq_names.index(sequence_name) 31 | # load valid frame range 32 | frames = np.loadtxt(self.range_files[index], dtype=int, delimiter=',') 33 | img_files = [os.path.join(self.seq_dirs[index], 'img/%04d.jpg' % f) for f in range(frames[0], frames[1] + 1)] 34 | 35 | # load annotations 36 | anno = np.loadtxt(self.anno_files[index], delimiter=',') 37 | assert len(img_files) == len(anno) 38 | assert anno.shape[1] == 4 39 | 40 | # return img_files, anno 41 | return Sequence(sequence_name, img_files, 'tc128', anno.reshape(-1, 4)) 42 | 43 | def __len__(self): 44 | return len(self.seq_names) 45 | -------------------------------------------------------------------------------- /external/lib/test/evaluation/trackingnetdataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList 3 | import os 4 | from lib.test.utils.load_text import load_text 5 | 6 | 7 | class TrackingNetDataset(BaseDataset): 8 | """ TrackingNet test set. 9 | 10 | Publication: 11 | TrackingNet: A Large-Scale Dataset and Benchmark for Object Tracking in the Wild. 12 | Matthias Mueller,Adel Bibi, Silvio Giancola, Salman Al-Subaihi and Bernard Ghanem 13 | ECCV, 2018 14 | https://ivul.kaust.edu.sa/Documents/Publications/2018/TrackingNet%20A%20Large%20Scale%20Dataset%20and%20Benchmark%20for%20Object%20Tracking%20in%20the%20Wild.pdf 15 | 16 | Download the dataset using the toolkit https://github.com/SilvioGiancola/TrackingNet-devkit. 17 | """ 18 | def __init__(self): 19 | super().__init__() 20 | self.base_path = self.env_settings.trackingnet_path 21 | 22 | sets = 'TEST' 23 | if not isinstance(sets, (list, tuple)): 24 | if sets == 'TEST': 25 | sets = ['TEST'] 26 | elif sets == 'TRAIN': 27 | sets = ['TRAIN_{}'.format(i) for i in range(5)] 28 | 29 | self.sequence_list = self._list_sequences(self.base_path, sets) 30 | 31 | def get_sequence_list(self): 32 | return SequenceList([self._construct_sequence(set, seq_name) for set, seq_name in self.sequence_list]) 33 | 34 | def _construct_sequence(self, set, sequence_name): 35 | anno_path = '{}/{}/anno/{}.txt'.format(self.base_path, set, sequence_name) 36 | 37 | ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64, backend='numpy') 38 | 39 | frames_path = '{}/{}/frames/{}'.format(self.base_path, set, sequence_name) 40 | frame_list = [frame for frame in os.listdir(frames_path) if frame.endswith(".jpg")] 41 | frame_list.sort(key=lambda f: int(f[:-4])) 42 | frames_list = [os.path.join(frames_path, frame) for frame in frame_list] 43 | 44 | return Sequence(sequence_name, frames_list, 'trackingnet', ground_truth_rect.reshape(-1, 4)) 45 | 46 | def __len__(self): 47 | return len(self.sequence_list) 48 | 49 | def _list_sequences(self, root, set_ids): 50 | sequence_list = [] 51 | 52 | for s in set_ids: 53 | anno_dir = os.path.join(root, s, "anno") 54 | sequences_cur_set = [(s, os.path.splitext(f)[0]) for f in os.listdir(anno_dir) if f.endswith('.txt')] 55 | 56 | sequence_list += sequences_cur_set 57 | 58 | return sequence_list 59 | -------------------------------------------------------------------------------- /external/lib/test/parameter/unicorn_sot.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 2 | from lib.test.utils import TrackerParams 3 | import os 4 | from lib.test.evaluation.environment import env_settings 5 | 6 | """ Parameters for Unicorn-SOT """ 7 | 8 | def parameters(exp_name: str): 9 | params = TrackerParams() 10 | prj_dir = env_settings().prj_dir 11 | 12 | # search region setting 13 | params.exp_name = exp_name 14 | # Network checkpoint path 15 | params.checkpoint = os.path.join(prj_dir, "Unicorn_outputs/%s/latest_ckpt.pth"%exp_name) 16 | 17 | return params -------------------------------------------------------------------------------- /external/lib/test/parameter/unicorn_vos.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 2 | from lib.test.utils import TrackerParams 3 | import os 4 | from lib.test.evaluation.environment import env_settings 5 | 6 | """ Parameters for Unicorn-VOS """ 7 | 8 | def parameters(exp_name: str): 9 | params = TrackerParams() 10 | prj_dir = env_settings().prj_dir 11 | 12 | # search region setting 13 | params.exp_name = exp_name 14 | # Network checkpoint path 15 | params.checkpoint = os.path.join(prj_dir, "Unicorn_outputs/%s/latest_ckpt.pth"%exp_name) 16 | 17 | return params -------------------------------------------------------------------------------- /external/lib/test/tracker/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /external/lib/test/tracker/basetracker.py: -------------------------------------------------------------------------------- 1 | from _collections import OrderedDict 2 | 3 | 4 | class BaseTracker: 5 | """Base class for all trackers.""" 6 | 7 | def __init__(self, params): 8 | self.params = params 9 | self.visdom = None 10 | 11 | def predicts_segmentation_mask(self): 12 | return False 13 | 14 | def initialize(self, image, info: dict) -> dict: 15 | """Overload this function in your tracker. This should initialize the model.""" 16 | raise NotImplementedError 17 | 18 | def track(self, image, info: dict = None) -> dict: 19 | """Overload this function in your tracker. This should track in the frame and update the model.""" 20 | raise NotImplementedError 21 | 22 | def visdom_draw_tracking(self, image, box, segmentation=None): 23 | if isinstance(box, OrderedDict): 24 | box = [v for k, v in box.items()] 25 | else: 26 | box = (box,) 27 | if segmentation is None: 28 | self.visdom.register((image, *box), 'Tracking', 1, 'Tracking') 29 | else: 30 | self.visdom.register((image, *box, segmentation), 'Tracking', 1, 'Tracking') -------------------------------------------------------------------------------- /external/lib/test/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .params import TrackerParams, FeatureParams -------------------------------------------------------------------------------- /external/lib/test/utils/_init_paths.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import os.path as osp 6 | import sys 7 | 8 | 9 | def add_path(path): 10 | if path not in sys.path: 11 | sys.path.insert(0, path) 12 | 13 | 14 | this_dir = osp.dirname(__file__) 15 | 16 | prj_path = osp.join(this_dir, '..', '..', '..') 17 | add_path(prj_path) -------------------------------------------------------------------------------- /external/lib/test/utils/load_text.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | 4 | 5 | def load_text_numpy(path, delimiter, dtype): 6 | if isinstance(delimiter, (tuple, list)): 7 | for d in delimiter: 8 | try: 9 | ground_truth_rect = np.loadtxt(path, delimiter=d, dtype=dtype) 10 | return ground_truth_rect 11 | except: 12 | pass 13 | 14 | raise Exception('Could not read file {}'.format(path)) 15 | else: 16 | ground_truth_rect = np.loadtxt(path, delimiter=delimiter, dtype=dtype) 17 | return ground_truth_rect 18 | 19 | 20 | def load_text_pandas(path, delimiter, dtype): 21 | if isinstance(delimiter, (tuple, list)): 22 | for d in delimiter: 23 | try: 24 | ground_truth_rect = pd.read_csv(path, delimiter=d, header=None, dtype=dtype, na_filter=False, 25 | low_memory=False).values 26 | return ground_truth_rect 27 | except Exception as e: 28 | pass 29 | 30 | raise Exception('Could not read file {}'.format(path)) 31 | else: 32 | ground_truth_rect = pd.read_csv(path, delimiter=delimiter, header=None, dtype=dtype, na_filter=False, 33 | low_memory=False).values 34 | return ground_truth_rect 35 | 36 | 37 | def load_text(path, delimiter=' ', dtype=np.float32, backend='numpy'): 38 | if backend == 'numpy': 39 | return load_text_numpy(path, delimiter, dtype) 40 | elif backend == 'pandas': 41 | return load_text_pandas(path, delimiter, dtype) -------------------------------------------------------------------------------- /external/lib/test/utils/params.py: -------------------------------------------------------------------------------- 1 | from lib.utils import TensorList 2 | import random 3 | 4 | 5 | class TrackerParams: 6 | """Class for tracker parameters.""" 7 | def set_default_values(self, default_vals: dict): 8 | for name, val in default_vals.items(): 9 | if not hasattr(self, name): 10 | setattr(self, name, val) 11 | 12 | def get(self, name: str, *default): 13 | """Get a parameter value with the given name. If it does not exists, it return the default value given as a 14 | second argument or returns an error if no default value is given.""" 15 | if len(default) > 1: 16 | raise ValueError('Can only give one default value.') 17 | 18 | if not default: 19 | return getattr(self, name) 20 | 21 | return getattr(self, name, default[0]) 22 | 23 | def has(self, name: str): 24 | """Check if there exist a parameter with the given name.""" 25 | return hasattr(self, name) 26 | 27 | 28 | class FeatureParams: 29 | """Class for feature specific parameters""" 30 | def __init__(self, *args, **kwargs): 31 | if len(args) > 0: 32 | raise ValueError 33 | 34 | for name, val in kwargs.items(): 35 | if isinstance(val, list): 36 | setattr(self, name, TensorList(val)) 37 | else: 38 | setattr(self, name, val) 39 | 40 | 41 | def Choice(*args): 42 | """Can be used to sample random parameter values.""" 43 | return random.choice(args) -------------------------------------------------------------------------------- /external/lib/test/utils/transform_got10k.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import shutil 4 | import argparse 5 | import _init_paths 6 | from lib.test.evaluation.environment import env_settings 7 | 8 | 9 | def transform_got10k(tracker_name, cfg_name): 10 | env = env_settings() 11 | result_dir = env.results_path 12 | src_dir = os.path.join(result_dir, "%s/%s/got10k/" % (tracker_name, cfg_name)) 13 | dest_dir = os.path.join(result_dir, "%s/%s/got10k_submit/" % (tracker_name, cfg_name)) 14 | if not os.path.exists(dest_dir): 15 | os.makedirs(dest_dir) 16 | items = os.listdir(src_dir) 17 | for item in items: 18 | if "all" in item: 19 | continue 20 | src_path = os.path.join(src_dir, item) 21 | if "time" not in item: 22 | seq_name = item.replace(".txt", '') 23 | seq_dir = os.path.join(dest_dir, seq_name) 24 | if not os.path.exists(seq_dir): 25 | os.makedirs(seq_dir) 26 | new_item = item.replace(".txt", '_001.txt') 27 | dest_path = os.path.join(seq_dir, new_item) 28 | bbox_arr = np.loadtxt(src_path, dtype=np.int, delimiter='\t') 29 | np.savetxt(dest_path, bbox_arr, fmt='%d', delimiter=',') 30 | else: 31 | seq_name = item.replace("_time.txt", '') 32 | seq_dir = os.path.join(dest_dir, seq_name) 33 | if not os.path.exists(seq_dir): 34 | os.makedirs(seq_dir) 35 | dest_path = os.path.join(seq_dir, item) 36 | os.system("cp %s %s" % (src_path, dest_path)) 37 | # make zip archive 38 | shutil.make_archive(src_dir, "zip", src_dir) 39 | shutil.make_archive(dest_dir, "zip", dest_dir) 40 | # Remove the original files 41 | shutil.rmtree(src_dir) 42 | shutil.rmtree(dest_dir) 43 | 44 | 45 | if __name__ == "__main__": 46 | parser = argparse.ArgumentParser(description='transform got10k results.') 47 | parser.add_argument('--tracker_name', type=str, help='Name of tracking method.') 48 | parser.add_argument('--cfg_name', type=str, help='Name of config file.') 49 | 50 | args = parser.parse_args() 51 | transform_got10k(args.tracker_name, args.cfg_name) -------------------------------------------------------------------------------- /external/lib/test/utils/transform_trackingnet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import shutil 4 | import argparse 5 | import _init_paths 6 | from lib.test.evaluation.environment import env_settings 7 | 8 | 9 | def transform_trackingnet(tracker_name, cfg_name): 10 | env = env_settings() 11 | result_dir = env.results_path 12 | src_dir = os.path.join(result_dir, "%s/%s/trackingnet/" % (tracker_name, cfg_name)) 13 | dest_dir = os.path.join(result_dir, "%s/%s/trackingnet_submit/" % (tracker_name, cfg_name)) 14 | if not os.path.exists(dest_dir): 15 | os.makedirs(dest_dir) 16 | items = os.listdir(src_dir) 17 | for item in items: 18 | if "all" in item: 19 | continue 20 | if "time" not in item: 21 | src_path = os.path.join(src_dir, item) 22 | dest_path = os.path.join(dest_dir, item) 23 | bbox_arr = np.loadtxt(src_path, dtype=np.int, delimiter='\t') 24 | np.savetxt(dest_path, bbox_arr, fmt='%d', delimiter=',') 25 | # make zip archive 26 | shutil.make_archive(src_dir, "zip", src_dir) 27 | shutil.make_archive(dest_dir, "zip", dest_dir) 28 | # Remove the original files 29 | shutil.rmtree(src_dir) 30 | shutil.rmtree(dest_dir) 31 | 32 | 33 | if __name__ == "__main__": 34 | parser = argparse.ArgumentParser(description='transform trackingnet results.') 35 | parser.add_argument('--tracker_name', type=str, help='Name of tracking method.') 36 | parser.add_argument('--cfg_name', type=str, help='Name of config file.') 37 | 38 | args = parser.parse_args() 39 | transform_trackingnet(args.tracker_name, args.cfg_name) -------------------------------------------------------------------------------- /external/lib/train/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /external/lib/train/_init_paths.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import os.path as osp 6 | import sys 7 | 8 | 9 | def add_path(path): 10 | if path not in sys.path: 11 | sys.path.insert(0, path) 12 | 13 | 14 | this_dir = osp.dirname(__file__) 15 | 16 | prj_path = osp.join(this_dir, '../..') 17 | add_path(prj_path) 18 | -------------------------------------------------------------------------------- /external/lib/train/admin/__init__.py: -------------------------------------------------------------------------------- 1 | from .environment import env_settings, create_default_local_file_train 2 | from .stats import AverageMeter, StatValue 3 | -------------------------------------------------------------------------------- /external/lib/train/admin/environment.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import os 3 | from collections import OrderedDict 4 | 5 | 6 | def create_default_local_file_train(workspace_dir, data_dir): 7 | path = os.path.join(os.path.dirname(__file__), 'local.py') 8 | 9 | empty_str = '\'\'' 10 | default_settings = OrderedDict({ 11 | 'workspace_dir': workspace_dir, 12 | 'tensorboard_dir': os.path.join(workspace_dir, 'tensorboard'), # Directory for tensorboard files. 13 | 'pretrained_networks': os.path.join(workspace_dir, 'pretrained_networks'), 14 | 'coco_dir': os.path.join(data_dir, 'COCO2017'), 15 | 'davis_dir': os.path.join(data_dir, 'DAVIS'), 16 | 'youtubevos_dir': os.path.join(data_dir, 'ytbvos18'), 17 | }) 18 | 19 | comment = {'workspace_dir': 'Base directory for saving network checkpoints.', 20 | 'tensorboard_dir': 'Directory for tensorboard files.'} 21 | 22 | with open(path, 'w') as f: 23 | f.write('class EnvironmentSettings:\n') 24 | f.write(' def __init__(self):\n') 25 | 26 | for attr, attr_val in default_settings.items(): 27 | comment_str = None 28 | if attr in comment: 29 | comment_str = comment[attr] 30 | if comment_str is None: 31 | if attr_val == empty_str: 32 | f.write(' self.{} = {}\n'.format(attr, attr_val)) 33 | else: 34 | f.write(' self.{} = \'{}\'\n'.format(attr, attr_val)) 35 | else: 36 | f.write(' self.{} = \'{}\' # {}\n'.format(attr, attr_val, comment_str)) 37 | 38 | 39 | def env_settings(): 40 | env_module_name = 'lib.train.admin.local' 41 | try: 42 | env_module = importlib.import_module(env_module_name) 43 | return env_module.EnvironmentSettings() 44 | except: 45 | env_file = os.path.join(os.path.dirname(__file__), 'local.py') 46 | raise RuntimeError( 47 | 'YOU HAVE NOT SETUP YOUR local.py!!!\n Go to "{}" and set all the paths you need. Then try to run again.'.format( 48 | env_file)) 49 | -------------------------------------------------------------------------------- /external/lib/train/admin/local.py: -------------------------------------------------------------------------------- 1 | from unicorn.data import get_unicorn_datadir 2 | import os 3 | 4 | class EnvironmentSettings: 5 | def __init__(self): 6 | self.workspace_dir = os.path.join(get_unicorn_datadir(), '..') # Base directory for saving network checkpoints. 7 | self.coco_dir = os.path.join(get_unicorn_datadir(), 'COCO') 8 | # VOS 9 | self.davis_dir = os.path.join(get_unicorn_datadir(), 'DAVIS') 10 | self.davis16_dir = os.path.join(get_unicorn_datadir(), 'DAVIS') 11 | self.youtubevos_dir = os.path.join(get_unicorn_datadir(), 'ytbvos18') # youtube-vos 2018 val 12 | -------------------------------------------------------------------------------- /external/lib/train/admin/settings.py: -------------------------------------------------------------------------------- 1 | from lib.train.admin.environment import env_settings 2 | 3 | 4 | class Settings: 5 | """ Training settings, e.g. the paths to datasets and networks.""" 6 | def __init__(self): 7 | self.set_default() 8 | 9 | def set_default(self): 10 | self.env = env_settings() 11 | self.use_gpu = True 12 | 13 | 14 | -------------------------------------------------------------------------------- /external/lib/train/admin/stats.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class StatValue: 4 | def __init__(self): 5 | self.clear() 6 | 7 | def reset(self): 8 | self.val = 0 9 | 10 | def clear(self): 11 | self.reset() 12 | self.history = [] 13 | 14 | def update(self, val): 15 | self.val = val 16 | self.history.append(self.val) 17 | 18 | 19 | class AverageMeter(object): 20 | """Computes and stores the average and current value""" 21 | def __init__(self): 22 | self.clear() 23 | self.has_new_data = False 24 | 25 | def reset(self): 26 | self.avg = 0 27 | self.val = 0 28 | self.sum = 0 29 | self.count = 0 30 | 31 | def clear(self): 32 | self.reset() 33 | self.history = [] 34 | 35 | def update(self, val, n=1): 36 | self.val = val 37 | self.sum += val * n 38 | self.count += n 39 | self.avg = self.sum / self.count 40 | 41 | def new_epoch(self): 42 | if self.count > 0: 43 | self.history.append(self.avg) 44 | self.reset() 45 | self.has_new_data = True 46 | else: 47 | self.has_new_data = False 48 | 49 | 50 | def topk_accuracy(output, target, topk=(1,)): 51 | """Computes the precision@k for the specified values of k""" 52 | single_input = not isinstance(topk, (tuple, list)) 53 | if single_input: 54 | topk = (topk,) 55 | 56 | maxk = max(topk) 57 | batch_size = target.size(0) 58 | 59 | _, pred = output.topk(maxk, 1, True, True) 60 | pred = pred.t() 61 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 62 | 63 | res = [] 64 | for k in topk: 65 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)[0] 66 | res.append(correct_k * 100.0 / batch_size) 67 | 68 | if single_input: 69 | return res[0] 70 | 71 | return res 72 | -------------------------------------------------------------------------------- /external/lib/train/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .loader import LTRLoader 2 | from .image_loader import jpeg4py_loader, opencv_loader, jpeg4py_loader_w_failsafe, default_image_loader -------------------------------------------------------------------------------- /external/lib/train/data/bounding_box_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def rect_to_rel(bb, sz_norm=None): 5 | """Convert standard rectangular parametrization of the bounding box [x, y, w, h] 6 | to relative parametrization [cx/sw, cy/sh, log(w), log(h)], where [cx, cy] is the center coordinate. 7 | args: 8 | bb - N x 4 tensor of boxes. 9 | sz_norm - [N] x 2 tensor of value of [sw, sh] (optional). sw=w and sh=h if not given. 10 | """ 11 | 12 | c = bb[...,:2] + 0.5 * bb[...,2:] 13 | if sz_norm is None: 14 | c_rel = c / bb[...,2:] 15 | else: 16 | c_rel = c / sz_norm 17 | sz_rel = torch.log(bb[...,2:]) 18 | return torch.cat((c_rel, sz_rel), dim=-1) 19 | 20 | 21 | def rel_to_rect(bb, sz_norm=None): 22 | """Inverts the effect of rect_to_rel. See above.""" 23 | 24 | sz = torch.exp(bb[...,2:]) 25 | if sz_norm is None: 26 | c = bb[...,:2] * sz 27 | else: 28 | c = bb[...,:2] * sz_norm 29 | tl = c - 0.5 * sz 30 | return torch.cat((tl, sz), dim=-1) 31 | 32 | 33 | def masks_to_bboxes(mask, fmt='c'): 34 | 35 | """ Convert a mask tensor to one or more bounding boxes. 36 | Note: This function is a bit new, make sure it does what it says. /Andreas 37 | :param mask: Tensor of masks, shape = (..., H, W) 38 | :param fmt: bbox layout. 'c' => "center + size" or (x_center, y_center, width, height) 39 | 't' => "top left + size" or (x_left, y_top, width, height) 40 | 'v' => "vertices" or (x_left, y_top, x_right, y_bottom) 41 | :return: tensor containing a batch of bounding boxes, shape = (..., 4) 42 | """ 43 | batch_shape = mask.shape[:-2] 44 | mask = mask.reshape((-1, *mask.shape[-2:])) 45 | bboxes = [] 46 | 47 | for m in mask: 48 | mx = m.sum(dim=-2).nonzero() 49 | my = m.sum(dim=-1).nonzero() 50 | bb = [mx.min(), my.min(), mx.max(), my.max()] if (len(mx) > 0 and len(my) > 0) else [0, 0, 0, 0] 51 | bboxes.append(bb) 52 | 53 | bboxes = torch.tensor(bboxes, dtype=torch.float32, device=mask.device) 54 | bboxes = bboxes.reshape(batch_shape + (4,)) 55 | 56 | if fmt == 'v': 57 | return bboxes 58 | 59 | x1 = bboxes[..., :2] 60 | s = bboxes[..., 2:] - x1 + 1 61 | 62 | if fmt == 'c': 63 | return torch.cat((x1 + 0.5 * s, s), dim=-1) 64 | elif fmt == 't': 65 | return torch.cat((x1, s), dim=-1) 66 | 67 | raise ValueError("Undefined bounding box layout '%s'" % fmt) 68 | 69 | 70 | def masks_to_bboxes_multi(mask, ids, fmt='c'): 71 | assert mask.dim() == 2 72 | bboxes = [] 73 | 74 | for id in ids: 75 | mx = (mask == id).sum(dim=-2).nonzero() 76 | my = (mask == id).float().sum(dim=-1).nonzero() 77 | bb = [mx.min(), my.min(), mx.max(), my.max()] if (len(mx) > 0 and len(my) > 0) else [0, 0, 0, 0] 78 | 79 | bb = torch.tensor(bb, dtype=torch.float32, device=mask.device) 80 | 81 | x1 = bb[:2] 82 | s = bb[2:] - x1 + 1 83 | 84 | if fmt == 'v': 85 | pass 86 | elif fmt == 'c': 87 | bb = torch.cat((x1 + 0.5 * s, s), dim=-1) 88 | elif fmt == 't': 89 | bb = torch.cat((x1, s), dim=-1) 90 | else: 91 | raise ValueError("Undefined bounding box layout '%s'" % fmt) 92 | bboxes.append(bb) 93 | 94 | return bboxes 95 | -------------------------------------------------------------------------------- /external/lib/train/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .davis import Davis 2 | from .youtubevos import YouTubeVOS 3 | -------------------------------------------------------------------------------- /external/lib/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .tensor import TensorDict, TensorList 2 | -------------------------------------------------------------------------------- /external/lib/utils/box_ops.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision.ops.boxes import box_area 3 | import numpy as np 4 | 5 | 6 | def box_cxcywh_to_xyxy(x): 7 | x_c, y_c, w, h = x.unbind(-1) 8 | b = [(x_c - 0.5 * w), (y_c - 0.5 * h), 9 | (x_c + 0.5 * w), (y_c + 0.5 * h)] 10 | return torch.stack(b, dim=-1) 11 | 12 | 13 | def box_xywh_to_xyxy(x): 14 | x1, y1, w, h = x.unbind(-1) 15 | b = [x1, y1, x1 + w, y1 + h] 16 | return torch.stack(b, dim=-1) 17 | 18 | 19 | def box_xyxy_to_xywh(x): 20 | x1, y1, x2, y2 = x.unbind(-1) 21 | b = [x1, y1, x2 - x1, y2 - y1] 22 | return torch.stack(b, dim=-1) 23 | 24 | 25 | def box_xyxy_to_cxcywh(x): 26 | x0, y0, x1, y1 = x.unbind(-1) 27 | b = [(x0 + x1) / 2, (y0 + y1) / 2, 28 | (x1 - x0), (y1 - y0)] 29 | return torch.stack(b, dim=-1) 30 | 31 | 32 | # modified from torchvision to also return the union 33 | '''Note that this function only supports shape (N,4)''' 34 | 35 | 36 | def box_iou(boxes1, boxes2): 37 | """ 38 | 39 | :param boxes1: (N, 4) (x1,y1,x2,y2) 40 | :param boxes2: (N, 4) (x1,y1,x2,y2) 41 | :return: 42 | """ 43 | area1 = box_area(boxes1) # (N,) 44 | area2 = box_area(boxes2) # (N,) 45 | 46 | lt = torch.max(boxes1[:, :2], boxes2[:, :2]) # (N,2) 47 | rb = torch.min(boxes1[:, 2:], boxes2[:, 2:]) # (N,2) 48 | 49 | wh = (rb - lt).clamp(min=0) # (N,2) 50 | inter = wh[:, 0] * wh[:, 1] # (N,) 51 | 52 | union = area1 + area2 - inter 53 | 54 | iou = inter / union 55 | return iou, union 56 | 57 | 58 | '''Note that this implementation is different from DETR's''' 59 | 60 | 61 | def generalized_box_iou(boxes1, boxes2): 62 | """ 63 | Generalized IoU from https://giou.stanford.edu/ 64 | 65 | The boxes should be in [x0, y0, x1, y1] format 66 | 67 | boxes1: (N, 4) 68 | boxes2: (N, 4) 69 | """ 70 | # degenerate boxes gives inf / nan results 71 | # so do an early check 72 | # try: 73 | assert (boxes1[:, 2:] >= boxes1[:, :2]).all() 74 | assert (boxes2[:, 2:] >= boxes2[:, :2]).all() 75 | iou, union = box_iou(boxes1, boxes2) # (N,) 76 | 77 | lt = torch.min(boxes1[:, :2], boxes2[:, :2]) 78 | rb = torch.max(boxes1[:, 2:], boxes2[:, 2:]) 79 | 80 | wh = (rb - lt).clamp(min=0) # (N,2) 81 | area = wh[:, 0] * wh[:, 1] # (N,) 82 | 83 | return iou - (area - union) / area, iou 84 | 85 | 86 | def giou_loss(boxes1, boxes2): 87 | """ 88 | 89 | :param boxes1: (N, 4) (x1,y1,x2,y2) 90 | :param boxes2: (N, 4) (x1,y1,x2,y2) 91 | :return: 92 | """ 93 | giou, iou = generalized_box_iou(boxes1, boxes2) 94 | return (1 - giou).mean(), iou.mean() 95 | 96 | 97 | def clip_box(box: list, H, W, margin=0): 98 | x1, y1, w, h = box 99 | x2, y2 = x1 + w, y1 + h 100 | x1 = min(max(0, x1), W-margin) 101 | x2 = min(max(margin, x2), W) 102 | y1 = min(max(0, y1), H-margin) 103 | y2 = min(max(margin, y2), H) 104 | w = max(margin, x2-x1) 105 | h = max(margin, y2-y1) 106 | return [x1, y1, w, h] 107 | -------------------------------------------------------------------------------- /external/lib/utils/get_image_anno_map.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | import argparse 4 | 5 | 6 | def parse_args(): 7 | parser = argparse.ArgumentParser(description="basic") 8 | parser.add_argument("--coco_root", type=str, default="/data/sdb/coco_2017", help="path to the coco root dir", required=True) 9 | parser.add_argument("--set", type=str, default="train", choices=["train", "val"], required=True) 10 | return parser.parse_args() 11 | 12 | 13 | if __name__ == "__main__": 14 | """Build a mapping between image_id and instances""" 15 | args = parse_args() 16 | coco_root = args.coco_root 17 | json_path = os.path.join(coco_root, 'annotations/instances_%s2017.json' % args.set) 18 | with open(json_path) as f: 19 | data = json.load(f) 20 | images_list = data['images'] # length 118287 21 | anno_list = data['annotations'] # length 860001 22 | result = {} 23 | for ann in anno_list: 24 | img_id = ann["image_id"] 25 | if img_id not in result: 26 | result[img_id] = [ann] 27 | else: 28 | result[img_id].append(ann) 29 | """deal with images without instances belonging to the specific 80 classes""" 30 | """there are 1021 images without instances""" 31 | for img in images_list: 32 | id = img["id"] 33 | if id not in result: 34 | result[id] = [] 35 | """save results""" 36 | result_path = os.path.join(coco_root, 'annotations/instances_%s2017_image_anno.json' % args.set) 37 | with open(result_path, "w") as f_w: 38 | json.dump(result, f_w) 39 | -------------------------------------------------------------------------------- /external/lib/utils/helpers.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | # torch 3 | import torch 4 | from torch.autograd import Variable 5 | from torch.utils import data 6 | 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.nn.init as init 10 | import torch.utils.model_zoo as model_zoo 11 | from torchvision import models 12 | 13 | # general libs 14 | import cv2 15 | import matplotlib.pyplot as plt 16 | from PIL import Image 17 | import numpy as np 18 | import time 19 | import os 20 | import copy 21 | 22 | 23 | def ToCuda(xs): 24 | if torch.cuda.is_available(): 25 | if isinstance(xs, list) or isinstance(xs, tuple): 26 | return [x.cuda() for x in xs] 27 | else: 28 | return xs.cuda() 29 | else: 30 | return xs 31 | 32 | 33 | def pad_divide_by(in_list, d, in_size): 34 | out_list = [] 35 | h, w = in_size 36 | if h % d > 0: 37 | new_h = h + d - h % d 38 | else: 39 | new_h = h 40 | if w % d > 0: 41 | new_w = w + d - w % d 42 | else: 43 | new_w = w 44 | lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2) 45 | lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2) 46 | pad_array = (int(lw), int(uw), int(lh), int(uh)) 47 | for inp in in_list: 48 | out_list.append(F.pad(inp, pad_array)) 49 | return out_list, pad_array 50 | 51 | 52 | def overlay_davis(image, mask, colors=[255, 0, 0], cscale=2, alpha=0.4): 53 | """ Overlay segmentation on top of RGB image. from davis official""" 54 | # import skimage 55 | from scipy.ndimage.morphology import binary_erosion, binary_dilation 56 | 57 | colors = np.reshape(colors, (-1, 3)) 58 | # colors = np.atleast_2d(colors) * cscale 59 | 60 | im_overlay = image.copy() 61 | object_ids = np.unique(mask) 62 | 63 | for object_id in object_ids[1:]: 64 | # Overlay color on binary mask 65 | foreground = image * alpha + np.ones(image.shape) * (1 - alpha) * np.array(colors[object_id]) 66 | binary_mask = mask == object_id 67 | 68 | # Compose image 69 | im_overlay[binary_mask] = foreground[binary_mask] 70 | 71 | # countours = skimage.morphology.binary.binary_dilation(binary_mask) - binary_mask 72 | countours = binary_dilation(binary_mask) ^ binary_mask 73 | # countours = cv2.dilate(binary_mask, cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))) - binary_mask 74 | im_overlay[countours, :] = 0 75 | 76 | return im_overlay.astype(image.dtype) 77 | -------------------------------------------------------------------------------- /external/lib/utils/lmdb_utils.py: -------------------------------------------------------------------------------- 1 | import lmdb 2 | import numpy as np 3 | import cv2 4 | import json 5 | 6 | LMDB_ENVS = dict() 7 | LMDB_HANDLES = dict() 8 | LMDB_FILELISTS = dict() 9 | 10 | 11 | def get_lmdb_handle(name): 12 | global LMDB_HANDLES, LMDB_FILELISTS 13 | item = LMDB_HANDLES.get(name, None) 14 | if item is None: 15 | env = lmdb.open(name, readonly=True, lock=False, readahead=False, meminit=False) 16 | LMDB_ENVS[name] = env 17 | item = env.begin(write=False) 18 | LMDB_HANDLES[name] = item 19 | 20 | return item 21 | 22 | 23 | def decode_img(lmdb_fname, key_name): 24 | handle = get_lmdb_handle(lmdb_fname) 25 | binfile = handle.get(key_name.encode()) 26 | if binfile is None: 27 | print("Illegal data detected. %s %s" % (lmdb_fname, key_name)) 28 | s = np.frombuffer(binfile, np.uint8) 29 | x = cv2.cvtColor(cv2.imdecode(s, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) 30 | return x 31 | 32 | 33 | def decode_str(lmdb_fname, key_name): 34 | handle = get_lmdb_handle(lmdb_fname) 35 | binfile = handle.get(key_name.encode()) 36 | string = binfile.decode() 37 | return string 38 | 39 | 40 | def decode_json(lmdb_fname, key_name): 41 | return json.loads(decode_str(lmdb_fname, key_name)) 42 | 43 | 44 | if __name__ == "__main__": 45 | lmdb_fname = "/data/sda/v-yanbi/iccv21/LittleBoy_clean/data/got10k_lmdb" 46 | '''Decode image''' 47 | # key_name = "test/GOT-10k_Test_000001/00000001.jpg" 48 | # img = decode_img(lmdb_fname, key_name) 49 | # cv2.imwrite("001.jpg", img) 50 | '''Decode str''' 51 | # key_name = "test/list.txt" 52 | # key_name = "train/GOT-10k_Train_000001/groundtruth.txt" 53 | key_name = "train/GOT-10k_Train_000001/absence.label" 54 | str_ = decode_str(lmdb_fname, key_name) 55 | print(str_) 56 | -------------------------------------------------------------------------------- /external/lib/utils/merge.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as NNF 3 | 4 | 5 | def merge_backbone_output(inp_list): 6 | """ merge outputs from backbone 7 | feat: HWxBxC, pos: HWxBxC, mask: BxHW 8 | """ 9 | nf = len(inp_list) 10 | seq_dict = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0), 11 | "mask": torch.cat([x["mask"] for x in inp_list], dim=1), 12 | "pos": torch.cat([x["pos"] for x in inp_list], dim=0), 13 | "h": inp_list[0]["h"], "w": inp_list[0]["w"], "nf": nf} 14 | return seq_dict 15 | 16 | 17 | def convert_to_onehot(x, dim): 18 | """convert to one-hot""" 19 | out_lbs = torch.zeros_like(x, device=x.device) 20 | pos_index = torch.argmax(x, dim=dim, keepdim=True) 21 | out_lbs.scatter_(dim, pos_index, 1) 22 | return out_lbs 23 | 24 | def adjust_labels_sz(inp_lbs, dh, dw): 25 | """ inp_lbs: PyTorch tensor (F, K, H, W) --> (F, K, H*s, W*s) 26 | dh: destination h, dw: destination w""" 27 | assert len(inp_lbs.size()) == 4 28 | """interpolation""" 29 | x = NNF.interpolate(inp_lbs, size=(dh, dw), mode="bilinear", align_corners=False) 30 | """convert to one-hot""" 31 | out_lbs = convert_to_onehot(x, dim=1) 32 | return out_lbs 33 | -------------------------------------------------------------------------------- /external/qdtrack/.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | mmdet/version.py 107 | data/ 108 | data 109 | .vscode 110 | .idea 111 | .DS_Store 112 | 113 | # custom 114 | *.pkl 115 | *.pkl.json 116 | *.log.json 117 | work_dirs/ 118 | 119 | # Pytorch 120 | *.pth 121 | *.py~ 122 | *.sh~ 123 | ckpts 124 | 125 | # results 126 | pcan/seg_track/ 127 | seg_track/ -------------------------------------------------------------------------------- /external/qdtrack/.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://gitlab.com/pycqa/flake8.git 3 | rev: 3.8.3 4 | hooks: 5 | - id: flake8 6 | - repo: https://github.com/asottile/seed-isort-config 7 | rev: v2.2.0 8 | hooks: 9 | - id: seed-isort-config 10 | - repo: https://github.com/timothycrosley/isort 11 | rev: 4.3.21 12 | hooks: 13 | - id: isort 14 | - repo: https://github.com/pre-commit/mirrors-yapf 15 | rev: v0.30.0 16 | hooks: 17 | - id: yapf 18 | - repo: https://github.com/pre-commit/pre-commit-hooks 19 | rev: v3.1.0 20 | hooks: 21 | - id: trailing-whitespace 22 | - id: check-yaml 23 | - id: end-of-file-fixer 24 | - id: requirements-txt-fixer 25 | - id: double-quote-string-fixer 26 | - id: check-merge-conflict 27 | - id: fix-encoding-pragma 28 | args: ["--remove"] 29 | - id: mixed-line-ending 30 | args: ["--fix=lf"] 31 | - repo: https://github.com/myint/docformatter 32 | rev: v1.3.1 33 | hooks: 34 | - id: docformatter 35 | args: ["--in-place", "--wrap-descriptions", "79"] 36 | -------------------------------------------------------------------------------- /external/qdtrack/clean_seg_track_json.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | 5 | 6 | if __name__ == "__main__": 7 | data_root = "data/bdd/labels/seg_track_20/polygons/train" 8 | file_list = os.listdir(data_root) 9 | for fname in file_list: 10 | if not fname.endswith(".json"): 11 | continue 12 | with open(os.path.join(data_root, fname), "r") as f: 13 | data = json.load(f) 14 | if "/" in data[0]["name"]: 15 | print(fname) 16 | for d in data: 17 | d["name"] = d["name"].split("/")[-1] 18 | with open(os.path.join(data_root, fname), "w") as f: 19 | json.dump(data, f) 20 | else: 21 | continue -------------------------------------------------------------------------------- /external/qdtrack/configs/_base_/qdtrack_faster_rcnn_r50_fpn.py: -------------------------------------------------------------------------------- 1 | _base_ = './faster_rcnn_r50_fpn.py' 2 | model = dict( 3 | type='QDTrack', 4 | rpn_head=dict( 5 | loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)), 6 | roi_head=dict( 7 | type='QuasiDenseRoIHead', 8 | track_roi_extractor=dict( 9 | type='SingleRoIExtractor', 10 | roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0), 11 | out_channels=256, 12 | featmap_strides=[4, 8, 16, 32]), 13 | track_head=dict( 14 | type='QuasiDenseEmbedHead', 15 | num_convs=4, 16 | num_fcs=1, 17 | embed_channels=256, 18 | norm_cfg=dict(type='GN', num_groups=32), 19 | loss_track=dict(type='MultiPosCrossEntropyLoss', loss_weight=0.25), 20 | loss_track_aux=dict( 21 | type='L2Loss', 22 | neg_pos_ub=3, 23 | pos_margin=0, 24 | neg_margin=0.1, 25 | hard_mining=True, 26 | loss_weight=1.0))), 27 | train_cfg=dict( 28 | embed=dict( 29 | assigner=dict( 30 | type='MaxIoUAssigner', 31 | pos_iou_thr=0.7, 32 | neg_iou_thr=0.3, 33 | min_pos_iou=0.5, 34 | match_low_quality=False, 35 | ignore_iof_thr=-1), 36 | sampler=dict( 37 | type='CombinedSampler', 38 | num=256, 39 | pos_fraction=0.5, 40 | neg_pos_ub=3, 41 | add_gt_as_proposals=True, 42 | pos_sampler=dict(type='InstanceBalancedPosSampler'), 43 | neg_sampler=dict(type='RandomSampler'))))) -------------------------------------------------------------------------------- /external/qdtrack/configs/bdd100k_mots/qdtrack-frcnn_r50_fpn_12e_bdd100k_evalseg.py: -------------------------------------------------------------------------------- 1 | _base_ = './qdtrack-frcnn_r50_fpn_12e_bdd100k.py' 2 | # dataset settings 3 | dataset_type = 'BDDVideoDataset' 4 | data_root = 'data/bdd/' 5 | img_norm_cfg = dict( 6 | mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) 7 | train_pipeline = [ 8 | dict(type='LoadMultiImagesFromFile'), 9 | dict(type='SeqLoadAnnotations', with_bbox=True, with_ins_id=True), 10 | dict(type='SeqResize', img_scale=(1296, 720), keep_ratio=True), 11 | dict(type='SeqRandomFlip', share_params=True, flip_ratio=0.5), 12 | dict(type='SeqNormalize', **img_norm_cfg), 13 | dict(type='SeqPad', size_divisor=32), 14 | dict(type='SeqDefaultFormatBundle'), 15 | dict( 16 | type='SeqCollect', 17 | keys=['img', 'gt_bboxes', 'gt_labels', 'gt_match_indices'], 18 | ref_prefix='ref'), 19 | ] 20 | test_pipeline = [ 21 | dict(type='LoadImageFromFile'), 22 | dict( 23 | type='MultiScaleFlipAug', 24 | img_scale=(1296, 720), 25 | flip=False, 26 | transforms=[ 27 | dict(type='Resize', keep_ratio=True), 28 | dict(type='RandomFlip'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | dict(type='Pad', size_divisor=32), 31 | dict(type='ImageToTensor', keys=['img']), 32 | dict(type='VideoCollect', keys=['img']) 33 | ]) 34 | ] 35 | data = dict( 36 | samples_per_gpu=2, 37 | workers_per_gpu=2, 38 | train=[ 39 | dict( 40 | type=dataset_type, 41 | ann_file=data_root + 'labels/seg_track_20/seg_track_train_cocoformat_new.json', 42 | img_prefix=data_root + 'images/seg_track_20/train', 43 | key_img_sampler=dict(interval=1), 44 | ref_img_sampler=dict(num_ref_imgs=1, scope=3, method='uniform'), 45 | pipeline=train_pipeline), 46 | dict( 47 | type=dataset_type, 48 | load_as_video=False, 49 | ann_file=data_root + 'labels/ins_seg/polygons/ins_seg_train_cocoformat.json', 50 | img_prefix=data_root + 'images/10k/train', 51 | pipeline=train_pipeline) 52 | ], 53 | val=dict( 54 | type=dataset_type, 55 | ann_file=data_root + 'labels/seg_track_20/seg_track_val_cocoformat.json', 56 | img_prefix=data_root + 'images/seg_track_20/val', 57 | pipeline=test_pipeline), 58 | test=dict( 59 | type=dataset_type, 60 | ann_file=data_root + 'labels/seg_track_20/seg_track_val_cocoformat.json', 61 | img_prefix=data_root + 'images/seg_track_20/val', 62 | pipeline=test_pipeline)) -------------------------------------------------------------------------------- /external/qdtrack/configs/bdd100k_mots/segtrack-frcnn_r50_fpn_12e_bdd10k_fixed.py: -------------------------------------------------------------------------------- 1 | _base_ = './segtrack-frcnn_r50_fpn_12e_bdd10k.py' 2 | 3 | model = dict(fixed=True) 4 | -------------------------------------------------------------------------------- /external/qdtrack/configs/bdd100k_mots/segtrack-frcnn_r50_fpn_12e_bdd10k_fixed_pcan.py: -------------------------------------------------------------------------------- 1 | _base_ = './segtrack-frcnn_r50_fpn_12e_bdd10k_fixed.py' 2 | # model settings 3 | model = dict( 4 | type='EMQuasiDenseMaskRCNNRefine', 5 | roi_head=dict( 6 | type='QuasiDenseSegRoIHeadRefine', 7 | double_train=False, 8 | mask_head=dict(type='FCNMaskHeadPlus'), 9 | refine_head=dict( 10 | type='EMMatchHeadPlus', 11 | num_convs=4, 12 | in_channels=256, 13 | conv_kernel_size=3, 14 | conv_out_channels=256, 15 | upsample_method='deconv', 16 | upsample_ratio=2, 17 | num_classes=8, 18 | pos_proto_num=10, #10 19 | neg_proto_num=10, #10 20 | stage_num=6, 21 | conv_cfg=None, 22 | norm_cfg=None, 23 | mask_thr_binary=0.5, 24 | match_score_thr=0.5, 25 | with_mask_ref=False, 26 | with_mask_key=True, 27 | with_dilation=False, 28 | loss_mask=dict( 29 | type='CrossEntropyLoss', 30 | use_sigmoid=True, 31 | loss_weight=1.0))), 32 | tracker=dict( 33 | type='QuasiDenseEmbedTracker', 34 | init_score_thr=0.5, 35 | obj_score_thr=0.3, 36 | match_score_thr=0.5, 37 | memo_tracklet_frames=10, 38 | memo_backdrop_frames=1, 39 | memo_momentum=1.0, 40 | nms_conf_thr=0.5, 41 | nms_backdrop_iou_thr=0.3, 42 | nms_class_iou_thr=0.7, 43 | with_cats=True, 44 | match_metric='bisoftmax'), 45 | ) 46 | 47 | load_from = './ckpts/segtrack-fixed-new.pth' 48 | -------------------------------------------------------------------------------- /external/qdtrack/configs/mot17/README.md: -------------------------------------------------------------------------------- 1 | ## Experiments on MOT17 Dataset 2 | 3 | ### 1. Download and prepare the data 4 | 5 | Please download the data from MOT Challenge. 6 | 7 | It is recommended to symlink the dataset root to `$QDTrack/data`. 8 | 9 | If your folder structure is different, you may need to change the corresponding paths in config files. 10 | 11 | Our folder structure follows 12 | 13 | ``` 14 | ├── qdtrack 15 | ├── tools 16 | ├── configs 17 | ├── data 18 | ├── MOT17 19 | ├── train 20 | ├── test 21 | ``` 22 | 23 | 24 | ### 2. Generate our annotation files 25 | 26 | ```shell 27 | python tools/convert_datasets/mot2coco.py -i ./data/MOT17 -o data/MOT17/annotations --split-train --convert-det 28 | ``` 29 | 30 | ### 3. Train the model on MOT17 31 | 32 | ```shell 33 | sh ./tools/dist_train.sh ./configs/mot17/qdtrack_frcnn_r50_fpn_4e_mot17.py 8 34 | ``` 35 | 36 | The pretrained model from MSCOCO can be obtained from [Google Drive](https://drive.google.com/file/d/1xK4Gvtd_2OchAyRY5WcoMkqoE9UM4gFO/view?usp=sharing) and [Baidu Yun](https://pan.baidu.com/s/1F9rLqsHjOv9DyJ-guvM-Qw) (passcode: vcu1). 37 | 38 | -------------------------------------------------------------------------------- /external/qdtrack/configs/tao/qdtrack_frcnn_r101_fpn_12e_tao_ft.py: -------------------------------------------------------------------------------- 1 | # model settings 2 | _base_ = './qdtrack_frcnn_r101_fpn_24e_lvis.py' 3 | model = dict(freeze_detector=True) 4 | # dataset settings 5 | img_norm_cfg = dict( 6 | mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) 7 | train_pipeline = [ 8 | dict(type='LoadMultiImagesFromFile'), 9 | dict(type='SeqLoadAnnotations', with_bbox=True, with_ins_id=True), 10 | dict( 11 | type='SeqResize', 12 | img_scale=[(1333, 640), (1333, 672), (1333, 704), (1333, 736), 13 | (1333, 768), (1333, 800)], 14 | share_params=True, 15 | multiscale_mode='value', 16 | keep_ratio=True), 17 | dict(type='SeqRandomFlip', share_params=True, flip_ratio=0.5), 18 | dict(type='SeqNormalize', **img_norm_cfg), 19 | dict(type='SeqPad', size_divisor=32), 20 | dict(type='SeqDefaultFormatBundle'), 21 | dict( 22 | type='SeqCollect', 23 | keys=['img', 'gt_bboxes', 'gt_labels', 'gt_match_indices'], 24 | ref_prefix='ref'), 25 | ] 26 | dataset_type = 'TaoDataset' 27 | data = dict( 28 | samples_per_gpu=2, 29 | workers_per_gpu=2, 30 | train=dict( 31 | _delete_=True, 32 | type='ClassBalancedDataset', 33 | oversample_thr=1e-3, 34 | dataset=dict( 35 | type=dataset_type, 36 | classes='data/tao/annotations/tao_classes.txt', 37 | ann_file='data/tao/annotations/train_482_ours.json', 38 | img_prefix='data/tao/frames/', 39 | key_img_sampler=dict(interval=1), 40 | ref_img_sampler=dict(num_ref_imgs=1, scope=1, method='uniform'), 41 | pipeline=train_pipeline))) 42 | optimizer = dict(type='SGD', lr=0.002, momentum=0.9, weight_decay=0.0001) 43 | lr_config = dict( 44 | policy='step', 45 | warmup='linear', 46 | warmup_iters=1000, 47 | warmup_ratio=1.0 / 1000, 48 | step=[8, 11]) 49 | total_epochs = 12 50 | load_from = 'ckpts/tao/qdtrack_r101_fpn_24e_pretrain-lvis+coco_20210811_132004-dc63f8a0.pth' 51 | evaluation = dict(metric=['track'], start=1, interval=1) 52 | work_dir = './work_dirs/tao/qdtrack_frcnn_r101_fpn_24e_tao_ft' 53 | -------------------------------------------------------------------------------- /external/qdtrack/docs/INSTALL.md: -------------------------------------------------------------------------------- 1 | 2 | ## Installation 3 | 4 | ### Requirements 5 | - Linux 6 | - Python 3.6+ 7 | - PyTorch 1.3+ 8 | - CUDA 9.2+ (If you build PyTorch from source, CUDA 9.0 is also compatible) 9 | - NCCL 2+ 10 | - GCC 5+ 11 | - [mmcv](https://github.com/open-mmlab/mmcv) 12 | - [mmdetection](https://github.com/open-mmlab/mmdetection) 13 | 14 | ### Install QDTrack 15 | 16 | a. Create a conda virtual environment and activate it. 17 | ```shell 18 | conda create -n qdtrack python=3.7 -y 19 | conda activate qdtrack 20 | ``` 21 | 22 | b. Install PyTorch and torchvision following the [official instructions](https://pytorch.org/), e.g., 23 | 24 | ```shell 25 | conda install pytorch torchvision -c pytorch 26 | ``` 27 | 28 | Note: Make sure that your compilation CUDA version and runtime CUDA version match. 29 | You can check the supported CUDA version for precompiled packages on the [PyTorch website](https://pytorch.org/). 30 | 31 | c. Install mmcv and mmdetection. 32 | 33 | ```shell 34 | pip install mmcv-full==1.3.10 35 | pip install mmdet==2.14.0 36 | ``` 37 | 38 | You can also refer to the [offical instructions](https://github.com/open-mmlab/mmdetection/blob/master/docs/install.md). 39 | 40 | d. Install QDTrack 41 | ```shell 42 | python setup.py develop 43 | ``` 44 | 45 | -------------------------------------------------------------------------------- /external/qdtrack/download_bdd100k.py: -------------------------------------------------------------------------------- 1 | import os 2 | import multiprocessing 3 | 4 | def download(url): 5 | os.system("wget -c %s" % url) 6 | 7 | if __name__ == "__main__": 8 | save_dir = "/opt/tiger/yanbin/BDD100K" 9 | if not os.path.exists(save_dir): 10 | os.makedirs(save_dir) 11 | url_list = [ 12 | "http://dl.yf.io/bdd100k/mots20/bdd100k_seg_track_20_images.zip", 13 | "http://dl.yf.io/bdd100k/mots20/bdd100k_seg_track_20_images.zip.md5", 14 | "http://dl.yf.io/bdd100k/mot20/images20-track-test-1.zip", 15 | "http://dl.yf.io/bdd100k/mot20/images20-track-test-1.zip.md5", 16 | "http://dl.yf.io/bdd100k/mot20/images20-track-test-2.zip", 17 | "http://dl.yf.io/bdd100k/mot20/images20-track-test-2.zip.md5", 18 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-1.zip", 19 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-1.zip.md5", 20 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-2.zip", 21 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-2.zip.md5", 22 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-3.zip", 23 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-3.zip.md5", 24 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-4.zip", 25 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-4.zip.md5", 26 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-5.zip", 27 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-5.zip.md5", 28 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-6.zip", 29 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-6.zip.md5", 30 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-7.zip", 31 | "http://dl.yf.io/bdd100k/mot20/images20-track-train-7.zip.md5", 32 | "http://dl.yf.io/bdd100k/mot20/images20-track-val-1.zip", 33 | "http://dl.yf.io/bdd100k/mot20/images20-track-val-1.zip.md5", 34 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_box_track_20_labels_trainval.zip", 35 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_seg_track_20_labels_trainval.zip", 36 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_images_100k.zip", 37 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_det_20_labels_trainval.zip", 38 | # instance segmentation 39 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_images_10k.zip", 40 | "https://bdd-data-storage-release.s3.us-west-2.amazonaws.com/bdd100k/2021/bdd100k_ins_seg_labels_trainval.zip", 41 | ] 42 | os.chdir(save_dir) 43 | for url in url_list: 44 | download(url) -------------------------------------------------------------------------------- /external/qdtrack/eval_bdd_submit.sh: -------------------------------------------------------------------------------- 1 | 2 | export PYTHONPATH=$PYTHONPATH:`pwd` 3 | 4 | python3 -m bdd100k.eval.run -t seg_track -g datasets/bdd/labels/seg_track_20/bitmasks/val -r seg_track -------------------------------------------------------------------------------- /external/qdtrack/figures/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/external/qdtrack/figures/teaser.png -------------------------------------------------------------------------------- /external/qdtrack/prepare_bdd100k.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | mkdir data 4 | 5 | pip3 install -U numpy 6 | 7 | ln -s /opt/tiger/yanbin/BDD100K/bdd100k /opt/tiger/Unicorn/external/qdtrack/data/bdd 8 | { 9 | python3 -m bdd100k.label.to_coco -m det -i data/bdd/labels/det_20/det_train.json -o data/bdd/labels/det_20/det_train_cocofmt.json 10 | } & 11 | { 12 | python3 -m bdd100k.label.to_coco -m det -i data/bdd/labels/det_20/det_val.json -o data/bdd/labels/det_20/det_val_cocofmt.json 13 | } & 14 | { 15 | python3 -m bdd100k.label.to_coco -m box_track -i data/bdd/labels/box_track_20/train -o data/bdd/labels/box_track_20/box_track_train_cocofmt.json 16 | } & 17 | { 18 | python3 -m bdd100k.label.to_coco -m box_track -i data/bdd/labels/box_track_20/val -o data/bdd/labels/box_track_20/box_track_val_cocofmt.json 19 | } & 20 | { 21 | python3 -m bdd100k.label.to_coco -m ins_seg -i data/bdd/labels/ins_seg/polygons/ins_seg_train.json -o data/bdd/labels/ins_seg/polygons/ins_seg_train_cocoformat.json -mb data/bdd/labels/ins_seg/bitmasks/train 22 | } & 23 | { 24 | python3 -m bdd100k.label.to_coco -m ins_seg -i data/bdd/labels/ins_seg/polygons/ins_seg_val.json -o data/bdd/labels/ins_seg/polygons/ins_seg_val_cocoformat.json -mb data/bdd/labels/ins_seg/bitmasks/val 25 | } & 26 | { 27 | python3 clean_seg_track_json.py 28 | python3 -m bdd100k.label.to_coco -m seg_track -i data/bdd/labels/seg_track_20/polygons/train -o data/bdd/labels/seg_track_20/seg_track_train_cocoformat.json -mb data/bdd/labels/seg_track_20/bitmasks/train 29 | } & 30 | { 31 | python3 -m bdd100k.label.to_coco -m seg_track -i data/bdd/labels/seg_track_20/polygons/val -o data/bdd/labels/seg_track_20/seg_track_val_cocoformat.json -mb data/bdd/labels/seg_track_20/bitmasks/val 32 | } & 33 | wait -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/VERSION: -------------------------------------------------------------------------------- 1 | 0.1.0 2 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import __version__, short_version 2 | 3 | __all__ = ['__version__', 'short_version'] 4 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/apis/__init__.py: -------------------------------------------------------------------------------- 1 | from .inference import inference_model, init_model 2 | from .test import multi_gpu_test, single_gpu_test 3 | from .test_omni import multi_gpu_test_omni 4 | from .train import train_model 5 | 6 | __all__ = [ 7 | 'init_model', 'inference_model', 'multi_gpu_test', 'single_gpu_test', 8 | 'train_model', 'multi_gpu_test_omni' 9 | ] 10 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/__init__.py: -------------------------------------------------------------------------------- 1 | from .evaluation import * # noqa: F401, F403 2 | from .track import * # noqa: F401, F403 3 | from .utils import * # noqa: F401, F403 4 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from .eval_hooks import EvalHook, DistEvalHook 2 | from .mot import eval_mot 3 | from .mot_pcan import xyxy2xywh 4 | 5 | __all__ = ['eval_mot', 'EvalHook', 'DistEvalHook', 'xyxy2xywh'] 6 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/evaluation/eval_hooks.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | 3 | import torch.distributed as dist 4 | from mmcv.runner import DistEvalHook as BaseDistEvalHook 5 | from mmcv.runner import EvalHook as BaseEvalHook 6 | from torch.nn.modules.batchnorm import _BatchNorm 7 | 8 | 9 | class EvalHook(BaseEvalHook): 10 | """Please refer to `mmcv.runner.hooks.evaluation.py:EvalHook` for detailed 11 | docstring.""" 12 | 13 | def _do_evaluate(self, runner): 14 | """perform evaluation and save ckpt.""" 15 | if not self._should_evaluate(runner): 16 | return 17 | 18 | if hasattr(self.dataloader.dataset, 19 | 'load_as_video') and self.dataloader.dataset.load_as_video: 20 | from qdtrack.apis import single_gpu_test 21 | else: 22 | from mmdet.apis import single_gpu_test 23 | results = single_gpu_test(runner.model, self.dataloader, show=False) 24 | runner.log_buffer.output['eval_iter_num'] = len(self.dataloader) 25 | key_score = self.evaluate(runner, results) 26 | if self.save_best: 27 | self._save_ckpt(runner, key_score) 28 | 29 | 30 | class DistEvalHook(BaseDistEvalHook): 31 | """Please refer to `mmcv.runner.hooks.evaluation.py:DistEvalHook` for 32 | detailed docstring.""" 33 | 34 | def _do_evaluate(self, runner): 35 | """perform evaluation and save ckpt.""" 36 | # Synchronization of BatchNorm's buffer (running_mean 37 | # and running_var) is not supported in the DDP of pytorch, 38 | # which may cause the inconsistent performance of models in 39 | # different ranks, so we broadcast BatchNorm's buffers 40 | # of rank 0 to other ranks to avoid this. 41 | if self.broadcast_bn_buffer: 42 | model = runner.model 43 | for name, module in model.named_modules(): 44 | if isinstance(module, 45 | _BatchNorm) and module.track_running_stats: 46 | dist.broadcast(module.running_var, 0) 47 | dist.broadcast(module.running_mean, 0) 48 | 49 | if not self._should_evaluate(runner): 50 | return 51 | 52 | tmpdir = self.tmpdir 53 | if tmpdir is None: 54 | tmpdir = osp.join(runner.work_dir, '.eval_hook') 55 | 56 | if hasattr(self.dataloader.dataset, 57 | 'load_as_video') and self.dataloader.dataset.load_as_video: 58 | from qdtrack.apis import multi_gpu_test 59 | else: 60 | from mmdet.apis import multi_gpu_test 61 | results = multi_gpu_test( 62 | runner.model, 63 | self.dataloader, 64 | tmpdir=tmpdir, 65 | gpu_collect=self.gpu_collect) 66 | if runner.rank == 0: 67 | print('\n') 68 | runner.log_buffer.output['eval_iter_num'] = len(self.dataloader) 69 | key_score = self.evaluate(runner, results) 70 | 71 | if self.save_best: 72 | self._save_ckpt(runner, key_score) 73 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/to_bdd100k/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import preds2bdd100k 2 | 3 | __all__ = ['preds2bdd100k'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/to_bdd100k/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | from functools import partial 4 | 5 | import pycocotools.mask as mask_utils 6 | import numpy as np 7 | import multiprocessing 8 | from multiprocessing import Pool 9 | from PIL import Image 10 | from tqdm import tqdm 11 | 12 | SHAPE = [720, 1280] 13 | 14 | 15 | def mask_prepare(track_dict): 16 | scores, colors, masks = [], [], [] 17 | for id_, instance in track_dict.items(): 18 | masks.append(mask_utils.decode(instance['segm'])) 19 | colors.append([instance['label'] + 1, 0, id_ >> 8, id_ & 255]) 20 | scores.append(instance['bbox'][-1]) 21 | return scores, colors, masks 22 | 23 | 24 | def mask_merge(mask_infor, img_name, bitmask_base): 25 | scores, colors, masks = mask_infor 26 | bitmask = np.zeros((*SHAPE, 4), dtype=np.uint8) 27 | sorted_idxs = np.argsort(scores) 28 | for idx in sorted_idxs: 29 | for i in range(4): 30 | bitmask[..., i] = ( 31 | bitmask[..., i] * (1 - masks[idx]) + 32 | masks[idx] * colors[idx][i]) 33 | bitmask_path = osp.join(bitmask_base, img_name.replace('.jpg', '.png')) 34 | bitmask_dir = osp.split(bitmask_path)[0] 35 | if not osp.exists(bitmask_dir): 36 | os.makedirs(bitmask_dir) 37 | bitmask = Image.fromarray(bitmask) 38 | bitmask.save(bitmask_path) 39 | 40 | 41 | def mask_merge_parallel(track_dicts, img_names, bitmask_base, nproc): 42 | # with Pool(nproc) as pool: 43 | print('\nCollecting mask information') 44 | mask_infors = [] 45 | for track_dict in track_dicts: 46 | mask_infors.append(mask_prepare(track_dict)) 47 | # mask_infors = pool.map(mask_prepare, tqdm(track_dicts)) 48 | print('\nMerging overlapped masks.') 49 | multiprocessing.set_start_method('spawn', force=True) 50 | with Pool(nproc) as pool: 51 | pool.starmap( 52 | partial(mask_merge, bitmask_base=bitmask_base), 53 | zip(mask_infors, img_names)) 54 | # for (mask_infor, img_name) in zip(mask_infors, img_names): 55 | # mask_merge(mask_infor, img_name, bitmask_base=bitmask_base) 56 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/track/__init__.py: -------------------------------------------------------------------------------- 1 | from .similarity import cal_similarity 2 | from .transforms import track2result, restore_result 3 | 4 | __all__ = ['cal_similarity', 'track2result', 'restore_result'] 5 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/track/similarity.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def cal_similarity(key_embeds, 6 | ref_embeds, 7 | method='dot_product', 8 | temperature=-1): 9 | assert method in ['dot_product', 'cosine'] 10 | 11 | if key_embeds.size(0) == 0 or ref_embeds.size(0) == 0: 12 | return torch.zeros((key_embeds.size(0), ref_embeds.size(0)), 13 | device=key_embeds.device) 14 | 15 | if method == 'cosine': 16 | key_embeds = F.normalize(key_embeds, p=2, dim=1) 17 | ref_embeds = F.normalize(ref_embeds, p=2, dim=1) 18 | return torch.mm(key_embeds, ref_embeds.t()) 19 | elif method == 'dot_product': 20 | if temperature > 0: 21 | dists = cal_similarity(key_embeds, ref_embeds, method='cosine') 22 | dists /= temperature 23 | return dists 24 | else: 25 | return torch.mm(key_embeds, ref_embeds.t()) 26 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/track/transforms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | 5 | def track2result(bboxes, labels, ids, num_classes): 6 | valid_inds = ids > -1 7 | bboxes = bboxes[valid_inds] 8 | labels = labels[valid_inds] 9 | ids = ids[valid_inds] 10 | 11 | if bboxes.shape[0] == 0: 12 | return [np.zeros((0, 6), dtype=np.float32) for i in range(num_classes)] 13 | else: 14 | if isinstance(bboxes, torch.Tensor): 15 | bboxes = bboxes.cpu().numpy() 16 | labels = labels.cpu().numpy() 17 | ids = ids.cpu().numpy() 18 | return [ 19 | np.concatenate((ids[labels == i, None], bboxes[labels == i, :]), 20 | axis=1) for i in range(num_classes) 21 | ] 22 | 23 | 24 | def restore_result(result, return_ids=False): 25 | labels = [] 26 | for i, bbox in enumerate(result): 27 | labels.extend([i] * bbox.shape[0]) 28 | bboxes = np.concatenate(result, axis=0).astype(np.float32) 29 | labels = np.array(labels, dtype=np.int64) 30 | if return_ids: 31 | ids = bboxes[:, 0].astype(np.int64) 32 | bboxes = bboxes[:, 1:] 33 | return bboxes, labels, ids 34 | else: 35 | return bboxes, labels 36 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/track/transforms_mots.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | 3 | 4 | def track2result(bboxes, labels, segms, ids): 5 | valid_inds = ids > -1 6 | bboxes = bboxes[valid_inds].cpu().numpy() 7 | labels = labels[valid_inds].cpu().numpy() 8 | segms = [segms[i] for i in range(len(segms)) if valid_inds[i] == True] 9 | ids = ids[valid_inds].cpu().numpy() 10 | 11 | outputs = defaultdict(list) 12 | for bbox, label, segm, id in zip(bboxes, labels, segms, ids): 13 | outputs[id] = dict(bbox=bbox, label=label, segm=segm) 14 | return outputs 15 | 16 | def segtrack2result(bboxes, labels, segms, ids): 17 | outputs = track2result(bboxes, labels, segms, ids) 18 | return outputs 19 | 20 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/core/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .visualization import imshow_mot_errors, imshow_tracks 2 | 3 | __all__ = ['imshow_tracks', 'imshow_mot_errors'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from mmdet.datasets.builder import (DATASETS, PIPELINES, build_dataset) 2 | 3 | from .bdd_video_dataset import BDDVideoDataset 4 | from .builder import build_dataloader 5 | from .coco_video_dataset import CocoVideoDataset 6 | from .parsers import CocoVID 7 | from .pipelines import (LoadMultiImagesFromFile, SeqCollect, 8 | SeqDefaultFormatBundle, SeqLoadAnnotations, 9 | SeqNormalize, SeqPad, SeqRandomFlip, SeqResize) 10 | from .tao_dataset import TaoDataset 11 | from .mot17_dataset import MOT17Dataset 12 | 13 | __all__ = [ 14 | 'DATASETS', 'PIPELINES', 'build_dataloader', 'build_dataset', 'CocoVID', 15 | 'BDDVideoDataset', 'CocoVideoDataset', 'LoadMultiImagesFromFile', 16 | 'SeqLoadAnnotations', 'SeqResize', 'SeqNormalize', 'SeqRandomFlip', 17 | 'SeqPad', 'SeqDefaultFormatBundle', 'SeqCollect', 'TaoDataset', 18 | 'MOT17Dataset' 19 | ] 20 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/bdd_video_dataset.py: -------------------------------------------------------------------------------- 1 | from mmdet.datasets import DATASETS 2 | 3 | from .coco_video_dataset import CocoVideoDataset 4 | 5 | 6 | @DATASETS.register_module() 7 | class BDDVideoDataset(CocoVideoDataset): 8 | 9 | CLASSES = ('pedestrian', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle') 10 | 11 | def __init__(self, *args, **kwargs): 12 | super().__init__(*args, **kwargs) 13 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/builder.py: -------------------------------------------------------------------------------- 1 | import random 2 | from functools import partial 3 | 4 | import numpy as np 5 | from mmcv.parallel import collate 6 | from mmcv.runner import get_dist_info 7 | from mmdet.datasets.samplers import DistributedGroupSampler, GroupSampler 8 | from torch.utils.data import DataLoader 9 | 10 | from .samplers import DistributedVideoSampler 11 | 12 | 13 | def build_dataloader(dataset, 14 | samples_per_gpu, 15 | workers_per_gpu, 16 | num_gpus=1, 17 | dist=True, 18 | shuffle=True, 19 | seed=None, 20 | **kwargs): 21 | """Build PyTorch DataLoader. 22 | 23 | In distributed training, each GPU/process has a dataloader. 24 | In non-distributed training, there is only one dataloader for all GPUs. 25 | 26 | Args: 27 | dataset (Dataset): A PyTorch dataset. 28 | samples_per_gpu (int): Number of training samples on each GPU, i.e., 29 | batch size of each GPU. 30 | workers_per_gpu (int): How many subprocesses to use for data loading 31 | for each GPU. 32 | num_gpus (int): Number of GPUs. Only used in non-distributed training. 33 | dist (bool): Distributed training/test or not. Default: True. 34 | shuffle (bool): Whether to shuffle the data at every epoch. 35 | Default: True. 36 | kwargs: any keyword argument to be used to initialize DataLoader 37 | 38 | Returns: 39 | DataLoader: A PyTorch dataloader. 40 | """ 41 | rank, world_size = get_dist_info() 42 | if dist: 43 | if shuffle: 44 | sampler = DistributedGroupSampler(dataset, samples_per_gpu, 45 | world_size, rank) 46 | else: 47 | sampler = DistributedVideoSampler( 48 | dataset, world_size, rank, shuffle=False) 49 | batch_size = samples_per_gpu 50 | num_workers = workers_per_gpu 51 | else: 52 | sampler = GroupSampler(dataset, samples_per_gpu) if shuffle else None 53 | batch_size = num_gpus * samples_per_gpu 54 | num_workers = num_gpus * workers_per_gpu 55 | 56 | init_fn = partial( 57 | worker_init_fn, num_workers=num_workers, rank=rank, 58 | seed=seed) if seed is not None else None 59 | 60 | data_loader = DataLoader( 61 | dataset, 62 | batch_size=batch_size, 63 | sampler=sampler, 64 | num_workers=num_workers, 65 | collate_fn=partial(collate, samples_per_gpu=samples_per_gpu), 66 | pin_memory=False, 67 | worker_init_fn=init_fn, 68 | **kwargs) 69 | 70 | return data_loader 71 | 72 | 73 | def worker_init_fn(worker_id, num_workers, rank, seed): 74 | # The seed of each worker equals to 75 | # num_worker * rank + worker_id + user_seed 76 | worker_seed = num_workers * rank + worker_id + seed 77 | np.random.seed(worker_seed) 78 | random.seed(worker_seed) 79 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/parsers/__init__.py: -------------------------------------------------------------------------------- 1 | from .coco_api import COCO, COCOeval 2 | from .coco_video_parser import CocoVID 3 | 4 | __all__ = ['COCO', 'COCOeval', 'CocoVID'] 5 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/parsers/coco_api.py: -------------------------------------------------------------------------------- 1 | # This file add snake case alias for coco api 2 | 3 | import warnings 4 | 5 | import pycocotools 6 | from pycocotools.coco import COCO as _COCO 7 | from pycocotools.cocoeval import COCOeval as _COCOeval 8 | 9 | 10 | class COCO(_COCO): 11 | """This class is almost the same as official pycocotools package. 12 | 13 | It implements some snake case function aliases. So that the COCO class has 14 | the same interface as LVIS class. 15 | """ 16 | 17 | def __init__(self, annotation_file=None): 18 | if getattr(pycocotools, '__version__', '0') >= '12.0.2': 19 | warnings.warn( 20 | 'mmpycocotools is deprecated. Please install official pycocotools by "pip install pycocotools"', # noqa: E501 21 | UserWarning) 22 | super().__init__(annotation_file=annotation_file) 23 | self.img_ann_map = self.imgToAnns 24 | self.cat_img_map = self.catToImgs 25 | 26 | def get_ann_ids(self, img_ids=[], cat_ids=[], area_rng=[], iscrowd=None): 27 | return self.getAnnIds(img_ids, cat_ids, area_rng, iscrowd) 28 | 29 | def get_cat_ids(self, cat_names=[], sup_names=[], cat_ids=[]): 30 | return self.getCatIds(cat_names, sup_names, cat_ids) 31 | 32 | def get_img_ids(self, img_ids=[], cat_ids=[]): 33 | return self.getImgIds(img_ids, cat_ids) 34 | 35 | def load_anns(self, ids): 36 | return self.loadAnns(ids) 37 | 38 | def load_cats(self, ids): 39 | return self.loadCats(ids) 40 | 41 | def load_imgs(self, ids): 42 | return self.loadImgs(ids) 43 | 44 | 45 | # just for the ease of import 46 | COCOeval = _COCOeval 47 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/parsers/coco_video_parser.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from collections import defaultdict 3 | from .coco_api import COCO 4 | from pycocotools.coco import _isArrayLike 5 | 6 | 7 | class CocoVID(COCO): 8 | 9 | def __init__(self, annotation_file=None): 10 | assert annotation_file, 'Annotation file must be provided.' 11 | super(CocoVID, self).__init__(annotation_file=annotation_file) 12 | 13 | def createIndex(self): 14 | print('creating index...') 15 | anns, cats, imgs, vids = {}, {}, {}, {} 16 | imgToAnns, catToImgs, vidToImgs = defaultdict(list), defaultdict( 17 | list), defaultdict(list) 18 | 19 | if 'videos' in self.dataset: 20 | for video in self.dataset['videos']: 21 | vids[video['id']] = video 22 | 23 | if 'annotations' in self.dataset: 24 | for ann in self.dataset['annotations']: 25 | imgToAnns[ann['image_id']].append(ann) 26 | anns[ann['id']] = ann 27 | 28 | if 'images' in self.dataset: 29 | for img in self.dataset['images']: 30 | vidToImgs[img['video_id']].append(img) 31 | imgs[img['id']] = img 32 | 33 | if 'categories' in self.dataset: 34 | for cat in self.dataset['categories']: 35 | cats[cat['id']] = cat 36 | 37 | if 'annotations' in self.dataset and 'categories' in self.dataset: 38 | for ann in self.dataset['annotations']: 39 | catToImgs[ann['category_id']].append(ann['image_id']) 40 | 41 | print('index created!') 42 | 43 | self.anns = anns 44 | self.imgToAnns = imgToAnns 45 | self.catToImgs = catToImgs 46 | self.imgs = imgs 47 | self.cats = cats 48 | self.videos = vids 49 | self.vidToImgs = vidToImgs 50 | 51 | def get_vid_ids(self, vidIds=[]): 52 | vidIds = vidIds if _isArrayLike(vidIds) else [vidIds] 53 | 54 | if len(vidIds) == 0: 55 | ids = self.videos.keys() 56 | else: 57 | ids = set(vidIds) 58 | 59 | return list(ids) 60 | 61 | def get_img_ids_from_vid(self, vidId): 62 | img_infos = self.vidToImgs[vidId] 63 | ids = list(np.zeros([len(img_infos)], dtype=np.int)) 64 | for img_info in img_infos: 65 | ids[img_info['frame_id']] = img_info['id'] 66 | return ids 67 | 68 | def load_vids(self, ids=[]): 69 | if _isArrayLike(ids): 70 | return [self.videos[id] for id in ids] 71 | elif type(ids) == int: 72 | return [self.videos[ids]] 73 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/pipelines/__init__.py: -------------------------------------------------------------------------------- 1 | from .formatting import VideoCollect, SeqCollect, SeqDefaultFormatBundle 2 | from .loading import LoadMultiImagesFromFile, SeqLoadAnnotations 3 | from .transforms import (SeqNormalize, SeqPad, SeqRandomFlip, SeqResize, 4 | SeqPhotoMetricDistortion, SeqRandomCrop) 5 | 6 | __all__ = [ 7 | 'LoadMultiImagesFromFile', 'SeqLoadAnnotations', 'SeqResize', 8 | 'SeqNormalize', 'SeqRandomFlip', 'SeqPad', 'SeqDefaultFormatBundle', 9 | 'SeqCollect', 'VideoCollect', 'SeqPhotoMetricDistortion', 'SeqRandomCrop' 10 | ] 11 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/pipelines/loading.py: -------------------------------------------------------------------------------- 1 | from mmdet.datasets.builder import PIPELINES 2 | from mmdet.datasets.pipelines import LoadAnnotations, LoadImageFromFile 3 | 4 | 5 | @PIPELINES.register_module() 6 | class LoadMultiImagesFromFile(LoadImageFromFile): 7 | 8 | def __init__(self, *args, **kwargs): 9 | super().__init__(*args, **kwargs) 10 | 11 | def __call__(self, results): 12 | outs = [] 13 | for _results in results: 14 | _results = super().__call__(_results) 15 | outs.append(_results) 16 | return outs 17 | 18 | 19 | @PIPELINES.register_module() 20 | class SeqLoadAnnotations(LoadAnnotations): 21 | 22 | def __init__(self, with_ins_id=False, *args, **kwargs): 23 | # TODO: name 24 | super().__init__(*args, **kwargs) 25 | self.with_ins_id = with_ins_id 26 | 27 | def _load_ins_ids(self, results): 28 | """Private function to load label annotations. 29 | 30 | Args: 31 | results (dict): Result dict from :obj:`mmdet.CustomDataset`. 32 | 33 | Returns: 34 | dict: The dict contains loaded label annotations. 35 | """ 36 | 37 | results['gt_match_indices'] = results['ann_info'][ 38 | 'match_indices'].copy() 39 | 40 | return results 41 | 42 | def __call__(self, results): 43 | outs = [] 44 | for _results in results: 45 | _results = super().__call__(_results) 46 | if self.with_ins_id: 47 | _results = self._load_ins_ids(_results) 48 | outs.append(_results) 49 | return outs 50 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | from .distributed_video_sampler import DistributedVideoSampler 2 | 3 | __all__ = ['DistributedVideoSampler'] 4 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/datasets/samplers/distributed_video_sampler.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch.utils.data import DistributedSampler as _DistributedSampler 3 | 4 | 5 | class DistributedVideoSampler(_DistributedSampler): 6 | 7 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False): 8 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) 9 | self.shuffle = shuffle 10 | assert not self.shuffle, 'Specific for video sequential testing.' 11 | self.num_samples = len(dataset) 12 | 13 | first_frame_indices = [] 14 | for i, img_info in enumerate(self.dataset.data_infos): 15 | if img_info['frame_id'] == 0: 16 | first_frame_indices.append(i) 17 | 18 | chunks = np.array_split(first_frame_indices, num_replicas) 19 | split_flags = [c[0] for c in chunks] 20 | split_flags.append(self.num_samples) 21 | 22 | self.indices = [ 23 | list(range(split_flags[i], split_flags[i + 1])) 24 | for i in range(self.num_replicas) 25 | ] 26 | 27 | def __iter__(self): 28 | indices = self.indices[self.rank] 29 | return iter(indices) 30 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import MODELS, TRACKERS, build_model, build_tracker 2 | from .losses import * # noqa: F401,F403 3 | from .mot import * # noqa: F401,F403 4 | from .roi_heads import * # noqa: F401,F403 5 | from .trackers import * # noqa: F401,F403 6 | 7 | __all__ = ['MODELS', 'TRACKERS', 'build_model', 'build_tracker'] 8 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry 2 | from mmcv.cnn import build_model_from_cfg as build 3 | 4 | MODELS = Registry('model') 5 | TRACKERS = Registry('tracker') 6 | 7 | 8 | def build_tracker(cfg): 9 | """Build tracker.""" 10 | return build(cfg, TRACKERS) 11 | 12 | 13 | def build_model(cfg, train_cfg=None, test_cfg=None): 14 | """Build model.""" 15 | return build(cfg, MODELS, dict(train_cfg=train_cfg, test_cfg=test_cfg)) 16 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .l2_loss import L2Loss 2 | from .multipos_cross_entropy_loss import MultiPosCrossEntropyLoss 3 | 4 | __all__ = ['L2Loss', 'MultiPosCrossEntropyLoss'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/losses/multipos_cross_entropy_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from mmdet.models import LOSSES, weight_reduce_loss 4 | 5 | 6 | def multi_pos_cross_entropy(pred, 7 | label, 8 | weight=None, 9 | reduction='mean', 10 | avg_factor=None): 11 | # element-wise losses 12 | # pos_inds = (label == 1).float() 13 | # neg_inds = (label == 0).float() 14 | # exp_pos = (torch.exp(-1 * pred) * pos_inds).sum(dim=1) 15 | # exp_neg = (torch.exp(pred.clamp(max=80)) * neg_inds).sum(dim=1) 16 | # loss = torch.log(1 + exp_pos * exp_neg) 17 | 18 | # a more numerical stable implementation. 19 | pos_inds = (label == 1) 20 | neg_inds = (label == 0) 21 | pred_pos = pred * pos_inds.float() 22 | pred_neg = pred * neg_inds.float() 23 | # use -inf to mask out unwanted elements. 24 | pred_pos[neg_inds] = pred_pos[neg_inds] + float('inf') 25 | pred_neg[pos_inds] = pred_neg[pos_inds] + float('-inf') 26 | 27 | _pos_expand = torch.repeat_interleave(pred_pos, pred.shape[1], dim=1) 28 | _neg_expand = pred_neg.repeat(1, pred.shape[1]) 29 | 30 | x = torch.nn.functional.pad((_neg_expand - _pos_expand), (0, 1), "constant", 0) 31 | loss = torch.logsumexp(x, dim=1) 32 | 33 | 34 | # apply weights and do the reduction 35 | if weight is not None: 36 | weight = weight.float() 37 | loss = weight_reduce_loss( 38 | loss, weight=weight, reduction=reduction, avg_factor=avg_factor) 39 | 40 | return loss 41 | 42 | 43 | @LOSSES.register_module() 44 | class MultiPosCrossEntropyLoss(nn.Module): 45 | 46 | def __init__(self, reduction='mean', loss_weight=1.0): 47 | super(MultiPosCrossEntropyLoss, self).__init__() 48 | self.reduction = reduction 49 | self.loss_weight = loss_weight 50 | 51 | def forward(self, 52 | cls_score, 53 | label, 54 | weight=None, 55 | avg_factor=None, 56 | reduction_override=None, 57 | **kwargs): 58 | assert cls_score.size() == label.size() 59 | assert reduction_override in (None, 'none', 'mean', 'sum') 60 | reduction = ( 61 | reduction_override if reduction_override else self.reduction) 62 | loss_cls = self.loss_weight * multi_pos_cross_entropy( 63 | cls_score, 64 | label, 65 | weight, 66 | reduction=reduction, 67 | avg_factor=avg_factor, 68 | **kwargs) 69 | return loss_cls 70 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/mot/__init__.py: -------------------------------------------------------------------------------- 1 | from .qdtrack import QDTrack 2 | 3 | __all__ = ['QDTrack'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/roi_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .quasi_dense_roi_head import QuasiDenseRoIHead 2 | from .track_heads import QuasiDenseEmbedHead 3 | 4 | __all__ = ['QuasiDenseRoIHead', 'QuasiDenseEmbedHead'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/roi_heads/track_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .quasi_dense_embed_head import QuasiDenseEmbedHead 2 | 3 | __all__ = ['QuasiDenseEmbedHead'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/models/trackers/__init__.py: -------------------------------------------------------------------------------- 1 | from .quasi_dense_embed_tracker import QuasiDenseEmbedTracker 2 | from .tao_tracker import TaoTracker 3 | 4 | __all__ = ['QuasiDenseEmbedTracker', 'TaoTracker'] -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .collect_env import collect_env 2 | from .logger import get_root_logger 3 | 4 | __all__ = ['collect_env', 'get_root_logger'] 5 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import subprocess 3 | import sys 4 | from collections import defaultdict 5 | 6 | import cv2 7 | import mmcv 8 | import torch 9 | import torchvision 10 | 11 | import qdtrack 12 | 13 | 14 | def collect_env(): 15 | env_info = {} 16 | env_info['sys.platform'] = sys.platform 17 | env_info['Python'] = sys.version.replace('\n', '') 18 | 19 | cuda_available = torch.cuda.is_available() 20 | env_info['CUDA available'] = cuda_available 21 | 22 | if cuda_available: 23 | from torch.utils.cpp_extension import CUDA_HOME 24 | env_info['CUDA_HOME'] = CUDA_HOME 25 | 26 | if CUDA_HOME is not None and osp.isdir(CUDA_HOME): 27 | try: 28 | nvcc = osp.join(CUDA_HOME, 'bin/nvcc') 29 | nvcc = subprocess.check_output( 30 | f'"{nvcc}" -V | tail -n1', shell=True) 31 | nvcc = nvcc.decode('utf-8').strip() 32 | except subprocess.SubprocessError: 33 | nvcc = 'Not Available' 34 | env_info['NVCC'] = nvcc 35 | 36 | devices = defaultdict(list) 37 | for k in range(torch.cuda.device_count()): 38 | devices[torch.cuda.get_device_name(k)].append(str(k)) 39 | for name, devids in devices.items(): 40 | env_info['GPU ' + ','.join(devids)] = name 41 | 42 | gcc = subprocess.check_output('gcc --version | head -n1', shell=True) 43 | gcc = gcc.decode('utf-8').strip() 44 | env_info['GCC'] = gcc 45 | 46 | env_info['PyTorch'] = torch.__version__ 47 | env_info['PyTorch compiling details'] = torch.__config__.show() 48 | 49 | env_info['TorchVision'] = torchvision.__version__ 50 | 51 | env_info['OpenCV'] = cv2.__version__ 52 | 53 | env_info['MMCV'] = mmcv.__version__ 54 | env_info['qdtrack'] = qdtrack.__version__ 55 | 56 | return env_info 57 | 58 | 59 | if __name__ == '__main__': 60 | for name, val in collect_env().items(): 61 | print(f'{name}: {val}') 62 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from mmcv.utils import get_logger 4 | 5 | 6 | def get_root_logger(log_file=None, log_level=logging.INFO): 7 | return get_logger('qdtrack', log_file, log_level) 8 | -------------------------------------------------------------------------------- /external/qdtrack/qdtrack/version.py: -------------------------------------------------------------------------------- 1 | # GENERATED VERSION FILE 2 | # TIME: Mon Aug 9 18:56:15 2021 3 | __version__ = '0.1.0+5532425' 4 | short_version = '0.1.0' 5 | version_info = (0, 1, 0) 6 | -------------------------------------------------------------------------------- /external/qdtrack/remap_cocofmt.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | from pathlib import Path 4 | import copy 5 | import argparse 6 | 7 | parser = argparse.ArgumentParser(description="remap category ids in coco-format box-track label files") 8 | parser.add_argument('--input', type=str) 9 | parser.add_argument('--output', type=str) 10 | args = parser.parse_args() 11 | 12 | path = Path(args.input) 13 | assert path.is_file(), f"{args.input} is not a valid path" 14 | new_path = Path(args.output) 15 | 16 | with open(path, 'r') as fh: 17 | config = json.load(fh) 18 | 19 | new_categories = [{'id': 1, 'name': 'pedestrian'}, {'id': 2, 'name': 'rider'}, {'id': 3, 'name': 'car'}, 20 | {'id': 4, 'name': 'bus'}, {'id': 5, 'name': 'truck'}, {'id': 6, 'name': 'bicycle'}, 21 | {'id': 7, 'name': 'motorcycle'}, {'id': 8, 'name': 'train'}] 22 | remap_cat_ids = { 23 | 1: 1, 24 | 2: 2, 25 | 3: 3, 26 | 4: 5, 27 | 5: 4, 28 | 6: 8, 29 | 7: 7, 30 | 8: 6 31 | } 32 | config["categories"] = new_categories 33 | remapped_annotations = [] 34 | for ann in config["annotations"]: 35 | c_ann = copy.deepcopy(ann) 36 | c_ann['category_id'] = remap_cat_ids[c_ann['category_id']] 37 | remapped_annotations.append(c_ann) 38 | 39 | config["annotations"] = remapped_annotations 40 | with open(str(new_path), 'w') as fh: 41 | json.dump(config, fh) -------------------------------------------------------------------------------- /external/qdtrack/requirements.txt: -------------------------------------------------------------------------------- 1 | lvis 2 | motmetrics 3 | numpy 4 | pycocotools 5 | torch>=1.1 6 | torchvision 7 | 8 | git+git://github.com/bdd100k/bdd100k.git 9 | -------------------------------------------------------------------------------- /external/qdtrack/setup.cfg: -------------------------------------------------------------------------------- 1 | [isort] 2 | line_length = 79 3 | multi_line_output = 0 4 | known_standard_library = setuptools 5 | known_first_party = qdtrack 6 | known_third_party = cv2,mmcv,mmdet,motmetrics,numpy,pandas,pycocotools,torch,torchvision,tqdm 7 | no_lines_before = STDLIB,LOCALFOLDER 8 | default_section = THIRDPARTY 9 | 10 | [yapf] 11 | BASED_ON_STYLE = pep8 12 | BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true 13 | SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true 14 | -------------------------------------------------------------------------------- /external/qdtrack/tools/debug_test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import mmcv 5 | import torch 6 | from mmcv import Config 7 | from mmdet.datasets import build_dataset 8 | from qdtrack.datasets import BDDVideoDataset 9 | from mmcv.runner import get_dist_info, init_dist 10 | import cv2 11 | from qdtrack.datasets import build_dataloader 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser(description='qdtrack test model') 15 | parser.add_argument('config', help='test config file path') 16 | parser.add_argument( 17 | '--launcher', 18 | choices=['none', 'pytorch', 'slurm', 'mpi'], 19 | default='none', 20 | help='job launcher') 21 | parser.add_argument('--local_rank', type=int, default=0) 22 | args = parser.parse_args() 23 | if 'LOCAL_RANK' not in os.environ: 24 | os.environ['LOCAL_RANK'] = str(args.local_rank) 25 | return args 26 | 27 | 28 | def main(): 29 | args = parse_args() 30 | 31 | cfg = Config.fromfile(args.config) 32 | 33 | # set cudnn_benchmark 34 | if cfg.get('cudnn_benchmark', False): 35 | torch.backends.cudnn.benchmark = True 36 | 37 | cfg.data.test.test_mode = True 38 | # init distributed env first, since logger depends on the dist info. 39 | if args.launcher == 'none': 40 | distributed = False 41 | else: 42 | distributed = True 43 | init_dist(args.launcher, **cfg.dist_params) 44 | # build the dataloader 45 | save_dir = "debug" 46 | if not os.path.exists(save_dir): 47 | os.makedirs(save_dir) 48 | dataset = build_dataset(cfg.data.test) 49 | data_loader = build_dataloader( 50 | dataset, 51 | samples_per_gpu=1, 52 | workers_per_gpu=cfg.data.workers_per_gpu, 53 | dist=distributed, 54 | shuffle=False) 55 | rank, _ = get_dist_info() 56 | for idx, data in enumerate(data_loader): 57 | if idx == 1: 58 | break 59 | print(rank, data) 60 | # img = data["img"][0].numpy().transpose(1, 2, 0) 61 | # cv2.imwrite(os.path.join(save_dir, "%02d.jpg" % idx), img) 62 | 63 | 64 | if __name__ == '__main__': 65 | main() 66 | -------------------------------------------------------------------------------- /external/qdtrack/tools/dist_debug_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | GPUS=$2 5 | PORT=${PORT:-29500} 6 | 7 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 8 | python3 -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 9 | $(dirname "$0")/debug_test.py $CONFIG $CHECKPOINT --launcher pytorch ${@:3} 10 | -------------------------------------------------------------------------------- /external/qdtrack/tools/dist_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | CHECKPOINT=$2 5 | GPUS=$3 6 | PORT=${PORT:-29500} 7 | 8 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 9 | python3 -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 10 | $(dirname "$0")/test.py $CONFIG $CHECKPOINT --launcher pytorch ${@:4} 11 | -------------------------------------------------------------------------------- /external/qdtrack/tools/dist_test_omni.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | CHECKPOINT=$2 5 | GPUS=$3 6 | EXP_NAME=$4 7 | PORT=${PORT:-29500} 8 | 9 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 10 | python3 -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 11 | $(dirname "$0")/test_omni.py $CONFIG $CHECKPOINT $EXP_NAME --launcher pytorch ${@:5} 12 | -------------------------------------------------------------------------------- /external/qdtrack/tools/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | GPUS=$2 5 | PORT=${PORT:-29500} 6 | 7 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 8 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 9 | $(dirname "$0")/train.py $CONFIG --launcher pytorch ${@:3} 10 | -------------------------------------------------------------------------------- /external/qdtrack/tools/slurm_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | CONFIG=$3 8 | CHECKPOINT=$4 9 | GPUS=${GPUS:-8} 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | PY_ARGS=${@:5} 13 | SRUN_ARGS=${SRUN_ARGS:-""} 14 | 15 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u tools/test.py ${CONFIG} ${CHECKPOINT} --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /external/qdtrack/tools/slurm_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | CONFIG=$3 8 | WORK_DIR=$4 9 | GPUS=${GPUS:-8} 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | SRUN_ARGS=${SRUN_ARGS:-""} 13 | PY_ARGS=${@:5} 14 | 15 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u tools/train.py ${CONFIG} --work-dir=${WORK_DIR} --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /external/qdtrack/tools/to_bdd100k.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import mmcv 5 | from mmcv import Config, DictAction 6 | from mmdet.datasets import build_dataset 7 | from qdtrack.core.to_bdd100k import preds2bdd100k 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser(description='qdtrack test model') 12 | parser.add_argument('config', help='test config file path') 13 | parser.add_argument('--res', help='output result file') 14 | parser.add_argument( 15 | '--bdd-dir', 16 | type=str, 17 | help='path to the folder that will contain files in bdd100k format') 18 | parser.add_argument( 19 | '--coco-file', 20 | type=str, 21 | help='path to that json file that is in COCO submission format') 22 | parser.add_argument( 23 | '--task', 24 | type=str, 25 | nargs='+', 26 | help='task types', 27 | choices=['det', 'ins_seg', 'box_track', 'seg_track']) 28 | parser.add_argument( 29 | '--nproc', 30 | type=int, 31 | help='number of process for mask merging') 32 | parser.add_argument( 33 | '--cfg-options', 34 | nargs='+', 35 | action=DictAction, 36 | help='override some settings in the used config, the key-value pair ' 37 | 'in xxx=yyy format will be merged into config file.') 38 | args = parser.parse_args() 39 | return args 40 | 41 | 42 | def main(): 43 | args = parse_args() 44 | 45 | if not os.path.isfile(args.res): 46 | raise ValueError('The result file does not exist.') 47 | 48 | cfg = Config.fromfile(args.config) 49 | 50 | if args.cfg_options is not None: 51 | cfg.merge_from_dict(args.cfg_options) 52 | 53 | if cfg.get('USE_MMDET', False): 54 | from mmdet.datasets import build_dataloader 55 | else: 56 | from qdtrack.datasets import build_dataloader 57 | 58 | # build the dataloader 59 | cfg.data.test.test_mode = True 60 | dataset = build_dataset(cfg.data.test) 61 | 62 | print(f'\nLoading results from {args.res}') 63 | results = mmcv.load(args.res) 64 | 65 | if args.coco_file: 66 | dataset.format_results(results, jsonfile_prefix=args.coco_file) 67 | if args.bdd_dir: 68 | print("converting results to bdd100k...") 69 | preds2bdd100k( 70 | dataset, results, args.task, out_base=args.bdd_dir, nproc=args.nproc) 71 | print("Conversion is done.") 72 | 73 | if __name__ == '__main__': 74 | main() 75 | -------------------------------------------------------------------------------- /launch_uni.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | import os 3 | import argparse 4 | import random 5 | import time 6 | 7 | """Unified launch script for object detection, instance segmentation and four tracking tasks""" 8 | def parse_args(): 9 | """ 10 | args for training. 11 | """ 12 | parser = argparse.ArgumentParser(description='Parse args for training') 13 | # for train 14 | parser.add_argument('--name', type=str, help='model name') 15 | parser.add_argument('--nproc_per_node', type=int, default=8, help="number of GPUs per node") # num_gpu per node 16 | parser.add_argument('--batch', type=int, help="global batch size") 17 | parser.add_argument('--mode', type=str, choices=["multiple", "distribute"], default="multiple", 18 | help="train on a single node or multiple nodes") 19 | parser.add_argument('--port', type=int, default=0) 20 | # The following args are required for "distributed" mode (training on multiple nodes) 21 | parser.add_argument('--master_address', type=str, help="address of the master node") 22 | parser.add_argument('--nnodes', type=int, help="the total number of nodes") 23 | parser.add_argument('--node_rank', type=int, help="the rank of the current node") 24 | 25 | args = parser.parse_args() 26 | 27 | return args 28 | 29 | 30 | def _get_rand_port(): 31 | hour = time.time() // 3600 32 | random.seed(int(hour)) 33 | return random.randrange(40000, 60000) 34 | 35 | 36 | def main(): 37 | args = parse_args() 38 | # change to current dir 39 | prj_dir = os.path.dirname(os.path.abspath(__file__)) 40 | os.chdir(prj_dir) 41 | # Get port 42 | if args.port > 0: 43 | master_port = args.port 44 | else: # This reduce the conflict possibility, but the port availablity is not guaranteed. 45 | master_port = _get_rand_port() 46 | # train 47 | file_name = "exps/default/%s" % args.name 48 | if args.mode == "multiple": 49 | train_cmd = "python3 tools/train.py -f %s -d %d -b %d -o --resume" % (file_name, args.nproc_per_node, args.batch) 50 | elif args.mode == "distribute": 51 | sub_cmd = "tools/train_dist.py -f %s -b %d -o --resume" % (file_name, args.batch) 52 | train_cmd = f'python3 -m torch.distributed.launch --nproc_per_node={args.nproc_per_node} \ 53 | --nnodes={args.nnodes} --node_rank={args.node_rank} --master_addr={args.master_address} --master_port={master_port} \ 54 | {sub_cmd}' 55 | else: 56 | raise ValueError("mode should be 'multiple' or 'distribute'.") 57 | os.system(train_cmd) 58 | 59 | 60 | if __name__ == "__main__": 61 | main() 62 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # TODO: Update with exact module version 2 | PyYAML 3 | cython 4 | cython_bbox 5 | pandas 6 | pycocotools 7 | git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI 8 | scipy 9 | colorama 10 | jpeg4py 11 | visdom 12 | git+https://github.com/rwightman/pytorch-image-models.git 13 | tb-nightly 14 | easydict 15 | tikzplotlib 16 | numpy 17 | torch>=1.7 18 | opencv_python 19 | loguru 20 | scikit-image 21 | tqdm 22 | torchvision 23 | Pillow 24 | thop 25 | ninja 26 | tabulate 27 | tensorboard 28 | lap 29 | motmetrics 30 | filterpy 31 | h5py 32 | prettytable 33 | # verified versions 34 | onnx==1.8.1 35 | onnxruntime==1.8.0 36 | onnx-simplifier==0.3.5 -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [isort] 2 | line_length = 100 3 | multi_line_output = 3 4 | balanced_wrapping = True 5 | known_standard_library = setuptools 6 | known_third_party = tqdm,loguru 7 | known_data_processing = cv2,numpy,scipy,PIL,matplotlib,scikit_image 8 | known_datasets = pycocotools 9 | known_deeplearning = torch,torchvision,caffe2,onnx,apex,timm,thop,torch2trt,tensorrt,openvino,onnxruntime 10 | known_myself = yolox 11 | sections = FUTURE,STDLIB,THIRDPARTY,data_processing,datasets,deeplearning,myself,FIRSTPARTY,LOCALFOLDER 12 | no_lines_before=STDLIB,THIRDPARTY,datasets 13 | default_section = FIRSTPARTY 14 | 15 | [flake8] 16 | max-line-length = 100 17 | max-complexity = 18 18 | exclude = __init__.py 19 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved 3 | 4 | import re 5 | import setuptools 6 | import glob 7 | from os import path 8 | import torch 9 | from torch.utils.cpp_extension import CppExtension 10 | 11 | torch_ver = [int(x) for x in torch.__version__.split(".")[:2]] 12 | assert torch_ver >= [1, 7], "Requires PyTorch >= 1.7" 13 | 14 | 15 | def get_extensions(): 16 | this_dir = path.dirname(path.abspath(__file__)) 17 | extensions_dir = path.join(this_dir, "unicorn", "layers", "csrc") 18 | 19 | main_source = path.join(extensions_dir, "vision.cpp") 20 | sources = glob.glob(path.join(extensions_dir, "**", "*.cpp")) 21 | 22 | sources = [main_source] + sources 23 | extension = CppExtension 24 | 25 | extra_compile_args = {"cxx": ["-O3"]} 26 | define_macros = [] 27 | 28 | include_dirs = [extensions_dir] 29 | 30 | ext_modules = [ 31 | extension( 32 | "unicorn._C", 33 | sources, 34 | include_dirs=include_dirs, 35 | define_macros=define_macros, 36 | extra_compile_args=extra_compile_args, 37 | ) 38 | ] 39 | 40 | return ext_modules 41 | 42 | 43 | with open("unicorn/__init__.py", "r") as f: 44 | version = re.search( 45 | r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]', 46 | f.read(), re.MULTILINE 47 | ).group(1) 48 | 49 | 50 | with open("README.md", "r", encoding="utf-8") as f: 51 | long_description = f.read() 52 | 53 | 54 | setuptools.setup( 55 | name="unicorn", 56 | version=version, 57 | author="basedet team", 58 | python_requires=">=3.6", 59 | long_description=long_description, 60 | ext_modules=get_extensions(), 61 | classifiers=["Programming Language :: Python :: 3", "Operating System :: OS Independent"], 62 | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, 63 | packages=setuptools.find_packages(), 64 | ) 65 | -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MasterBin-IIAU/Unicorn/4da907948e7473280c57208dd9c9b6274c488fe7/tools/__init__.py -------------------------------------------------------------------------------- /tools/_init_paths.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import os 6 | import sys 7 | 8 | 9 | def add_path(path): 10 | if path not in sys.path: 11 | sys.path.insert(0, path) 12 | 13 | 14 | 15 | external_dir = os.path.join(os.path.dirname(__file__), "..", "external") 16 | if external_dir not in sys.path: 17 | add_path(external_dir) 18 | -------------------------------------------------------------------------------- /tools/analysis_results.py: -------------------------------------------------------------------------------- 1 | import _init_paths 2 | import matplotlib.pyplot as plt 3 | plt.rcParams['figure.figsize'] = [8, 8] 4 | import argparse 5 | from lib.test.analysis.plot_results import plot_results, print_results, print_per_sequence_results 6 | from lib.test.evaluation import get_dataset, trackerlist 7 | 8 | def parse_args(): 9 | """ 10 | args for training. 11 | """ 12 | parser = argparse.ArgumentParser(description='Parse args for training') 13 | # for evaluation 14 | parser.add_argument('--name', type=str, help='model name') 15 | 16 | args = parser.parse_args() 17 | 18 | return args 19 | 20 | if __name__ == "__main__": 21 | trackers = [] 22 | dataset_name = 'LaSOT' 23 | 24 | args = parse_args() 25 | trackers.extend(trackerlist(name='unicorn_sot', parameter_name=args.name, dataset_name=dataset_name, 26 | run_ids=None, display_name='Unicorn')) 27 | 28 | dataset = get_dataset('lasot') 29 | # plot_results(trackers, dataset, 'OTB2015', merge_results=True, plot_types=('success', 'norm_prec'), 30 | # skip_missing_seq=False, force_evaluation=True, plot_bin_gap=0.05) 31 | print_results(trackers, dataset, dataset_name, merge_results=True, plot_types=('success', 'prec', 'norm_prec')) 32 | # print_per_sequence_results(trackers, dataset, "Unicorn") 33 | -------------------------------------------------------------------------------- /tools/convert_cityperson_to_coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import json 4 | from PIL import Image 5 | 6 | DATA_PATH = 'datasets/Cityscapes/' 7 | DATA_FILE_PATH = 'datasets/data_path/citypersons.train' 8 | OUT_PATH = DATA_PATH + 'annotations/' 9 | 10 | def load_paths(data_path): 11 | with open(data_path, 'r') as file: 12 | img_files = file.readlines() 13 | img_files = [x.replace('\n', '') for x in img_files] 14 | img_files = list(filter(lambda x: len(x) > 0, img_files)) 15 | label_files = [x.replace('images', 'labels_with_ids').replace('.png', '.txt').replace('.jpg', '.txt') for x in img_files] 16 | return img_files, label_files 17 | 18 | if __name__ == '__main__': 19 | if not os.path.exists(OUT_PATH): 20 | os.mkdir(OUT_PATH) 21 | 22 | out_path = OUT_PATH + 'train.json' 23 | out = {'images': [], 'annotations': [], 'categories': [{'id': 1, 'name': 'person'}]} 24 | img_paths, label_paths = load_paths(DATA_FILE_PATH) 25 | image_cnt = 0 26 | ann_cnt = 0 27 | video_cnt = 0 28 | for img_path, label_path in zip(img_paths, label_paths): 29 | image_cnt += 1 30 | im = Image.open(os.path.join("datasets", img_path)) 31 | image_info = {'file_name': img_path, 32 | 'id': image_cnt, 33 | 'height': im.size[1], 34 | 'width': im.size[0]} 35 | out['images'].append(image_info) 36 | # Load labels 37 | if os.path.isfile(os.path.join("datasets", label_path)): 38 | labels0 = np.loadtxt(os.path.join("datasets", label_path), dtype=np.float32).reshape(-1, 6) 39 | # Normalized xywh to pixel xyxy format 40 | labels = labels0.copy() 41 | labels[:, 2] = image_info['width'] * (labels0[:, 2] - labels0[:, 4] / 2) 42 | labels[:, 3] = image_info['height'] * (labels0[:, 3] - labels0[:, 5] / 2) 43 | labels[:, 4] = image_info['width'] * labels0[:, 4] 44 | labels[:, 5] = image_info['height'] * labels0[:, 5] 45 | else: 46 | labels = np.array([]) 47 | for i in range(len(labels)): 48 | ann_cnt += 1 49 | fbox = labels[i, 2:6].tolist() 50 | ann = {'id': ann_cnt, 51 | 'category_id': 1, 52 | 'image_id': image_cnt, 53 | 'track_id': -1, 54 | 'bbox': fbox, 55 | 'area': fbox[2] * fbox[3], 56 | 'iscrowd': 0} 57 | out['annotations'].append(ann) 58 | print('loaded train for {} images and {} samples'.format(len(out['images']), len(out['annotations']))) 59 | json.dump(out, open(out_path, 'w')) 60 | -------------------------------------------------------------------------------- /tools/convert_crowdhuman_to_coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import json 4 | from PIL import Image 5 | 6 | DATA_PATH = 'datasets/crowdhuman/' 7 | OUT_PATH = DATA_PATH + 'annotations/' 8 | SPLITS = ['val', 'train'] 9 | DEBUG = False 10 | 11 | def load_func(fpath): 12 | print('fpath', fpath) 13 | assert os.path.exists(fpath) 14 | with open(fpath,'r') as fid: 15 | lines = fid.readlines() 16 | records =[json.loads(line.strip('\n')) for line in lines] 17 | return records 18 | 19 | if __name__ == '__main__': 20 | if not os.path.exists(OUT_PATH): 21 | os.mkdir(OUT_PATH) 22 | for split in SPLITS: 23 | data_path = DATA_PATH + split 24 | out_path = OUT_PATH + '{}.json'.format(split) 25 | out = {'images': [], 'annotations': [], 'categories': [{'id': 1, 'name': 'person'}]} 26 | ann_path = DATA_PATH + 'annotation_{}.odgt'.format(split) 27 | anns_data = load_func(ann_path) 28 | image_cnt = 0 29 | ann_cnt = 0 30 | video_cnt = 0 31 | for ann_data in anns_data: 32 | image_cnt += 1 33 | file_path = DATA_PATH + 'CrowdHuman_{}/'.format(split) + '{}.jpg'.format(ann_data['ID']) 34 | im = Image.open(file_path) 35 | image_info = {'file_name': '{}.jpg'.format(ann_data['ID']), 36 | 'id': image_cnt, 37 | 'height': im.size[1], 38 | 'width': im.size[0]} 39 | out['images'].append(image_info) 40 | if split != 'test': 41 | anns = ann_data['gtboxes'] 42 | for i in range(len(anns)): 43 | ann_cnt += 1 44 | fbox = anns[i]['fbox'] 45 | ann = {'id': ann_cnt, 46 | 'category_id': 1, 47 | 'image_id': image_cnt, 48 | 'track_id': -1, 49 | 'bbox_vis': anns[i]['vbox'], 50 | 'bbox': fbox, 51 | 'area': fbox[2] * fbox[3], 52 | 'iscrowd': 1 if 'extra' in anns[i] and \ 53 | 'ignore' in anns[i]['extra'] and \ 54 | anns[i]['extra']['ignore'] == 1 else 0} 55 | out['annotations'].append(ann) 56 | print('loaded {} for {} images and {} samples'.format(split, len(out['images']), len(out['annotations']))) 57 | json.dump(out, open(out_path, 'w')) -------------------------------------------------------------------------------- /tools/convert_ethz_to_coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import json 4 | from PIL import Image 5 | 6 | DATA_PATH = 'datasets/ETHZ/' 7 | DATA_FILE_PATH = 'datasets/data_path/eth.train' 8 | OUT_PATH = DATA_PATH + 'annotations/' 9 | 10 | def load_paths(data_path): 11 | with open(data_path, 'r') as file: 12 | img_files = file.readlines() 13 | img_files = [x.replace('\n', '') for x in img_files] 14 | img_files = list(filter(lambda x: len(x) > 0, img_files)) 15 | label_files = [x.replace('images', 'labels_with_ids').replace('.png', '.txt').replace('.jpg', '.txt') for x in img_files] 16 | return img_files, label_files 17 | 18 | if __name__ == '__main__': 19 | if not os.path.exists(OUT_PATH): 20 | os.mkdir(OUT_PATH) 21 | 22 | out_path = OUT_PATH + 'train.json' 23 | out = {'images': [], 'annotations': [], 'categories': [{'id': 1, 'name': 'person'}]} 24 | img_paths, label_paths = load_paths(DATA_FILE_PATH) 25 | image_cnt = 0 26 | ann_cnt = 0 27 | video_cnt = 0 28 | for img_path, label_path in zip(img_paths, label_paths): 29 | image_cnt += 1 30 | im = Image.open(os.path.join("datasets", img_path)) 31 | image_info = {'file_name': img_path, 32 | 'id': image_cnt, 33 | 'height': im.size[1], 34 | 'width': im.size[0]} 35 | out['images'].append(image_info) 36 | # Load labels 37 | if os.path.isfile(os.path.join("datasets", label_path)): 38 | labels0 = np.loadtxt(os.path.join("datasets", label_path), dtype=np.float32).reshape(-1, 6) 39 | # Normalized xywh to pixel xyxy format 40 | labels = labels0.copy() 41 | labels[:, 2] = image_info['width'] * (labels0[:, 2] - labels0[:, 4] / 2) 42 | labels[:, 3] = image_info['height'] * (labels0[:, 3] - labels0[:, 5] / 2) 43 | labels[:, 4] = image_info['width'] * labels0[:, 4] 44 | labels[:, 5] = image_info['height'] * labels0[:, 5] 45 | else: 46 | labels = np.array([]) 47 | for i in range(len(labels)): 48 | ann_cnt += 1 49 | fbox = labels[i, 2:6].tolist() 50 | ann = {'id': ann_cnt, 51 | 'category_id': 1, 52 | 'image_id': image_cnt, 53 | 'track_id': -1, 54 | 'bbox': fbox, 55 | 'area': fbox[2] * fbox[3], 56 | 'iscrowd': 0} 57 | out['annotations'].append(ann) 58 | print('loaded train for {} images and {} samples'.format(len(out['images']), len(out['annotations']))) 59 | json.dump(out, open(out_path, 'w')) 60 | -------------------------------------------------------------------------------- /tools/convert_mot17_to_omni.py: -------------------------------------------------------------------------------- 1 | from unicorn.data import get_unicorn_datadir 2 | from unicorn.data.datasets import MOTDataset 3 | import os 4 | import json 5 | import argparse 6 | 7 | def parse_args(): 8 | """ 9 | args for training. 10 | """ 11 | parser = argparse.ArgumentParser(description='Parse args for training') 12 | # for train 13 | parser.add_argument('--dataset_name', type=str, default="mot") 14 | parser.add_argument('--ori_json_file', type=str, default="train.json") 15 | parser.add_argument('--new_json_file', type=str, default="train_omni.json") 16 | 17 | args = parser.parse_args() 18 | 19 | return args 20 | 21 | if __name__ == "__main__": 22 | args = parse_args() 23 | dataset_name = args.dataset_name 24 | ori_json_file = args.ori_json_file 25 | new_json_file = args.new_json_file 26 | data_dir = os.path.join(get_unicorn_datadir(), dataset_name) 27 | save_path = os.path.join(data_dir, "annotations", new_json_file) 28 | dataset = MOTDataset(data_dir=data_dir, json_file=ori_json_file) 29 | omni_json = {} 30 | for (res, img_info, file_name) in dataset.annotations: 31 | (height, width, frame_id, video_id, file_name) = img_info 32 | if video_id not in omni_json: 33 | omni_json[video_id] = {} 34 | omni_json[video_id][frame_id] = {"res": res.tolist(), "img_info": img_info, "file_name": file_name} 35 | json.dump(omni_json, open(save_path, "w")) 36 | -------------------------------------------------------------------------------- /tools/copy_1to3.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | if __name__ == "__main__": 5 | exp_name = "unicorn_track_large_mot_challenge" 6 | src_dir = "Unicorn_outputs/%s/track_results_train" 7 | des_dir = "Unicorn_outputs/%s/track_results_dti" 8 | if not os.path.exists(des_dir): 9 | os.makedirs(des_dir) 10 | file_list = os.listdir(src_dir) 11 | for f in file_list: 12 | src_path = os.path.join(src_dir, f) 13 | assert ("FRCNN" in src_path) 14 | des_path1 = src_path.replace(src_dir, des_dir) 15 | des_path2 = des_path1.replace("FRCNN", "DPM") 16 | des_path3 = des_path1.replace("FRCNN", "SDP") 17 | os.system("cp %s %s" %(src_path, des_path1)) 18 | os.system("cp %s %s" %(src_path, des_path2)) 19 | os.system("cp %s %s" %(src_path, des_path3)) 20 | -------------------------------------------------------------------------------- /tools/export_torchscript.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | import argparse 6 | import os 7 | from loguru import logger 8 | 9 | import torch 10 | from torch import nn 11 | 12 | from unicorn.exp import get_exp 13 | 14 | 15 | def make_parser(): 16 | parser = argparse.ArgumentParser("YOLOX torchscript deploy") 17 | parser.add_argument( 18 | "--output-name", type=str, default="unicorn.torchscript.pt", help="output name of models" 19 | ) 20 | parser.add_argument("--batch-size", type=int, default=1, help="batch size") 21 | parser.add_argument( 22 | "-f", 23 | "--exp_file", 24 | default=None, 25 | type=str, 26 | help="expriment description file", 27 | ) 28 | parser.add_argument("-expn", "--experiment-name", type=str, default=None) 29 | parser.add_argument("-n", "--name", type=str, default=None, help="model name") 30 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path") 31 | parser.add_argument( 32 | "opts", 33 | help="Modify config options using the command-line", 34 | default=None, 35 | nargs=argparse.REMAINDER, 36 | ) 37 | 38 | return parser 39 | 40 | 41 | @logger.catch 42 | def main(): 43 | args = make_parser().parse_args() 44 | logger.info("args value: {}".format(args)) 45 | exp = get_exp(args.exp_file, args.name) 46 | exp.merge(args.opts) 47 | 48 | if not args.experiment_name: 49 | args.experiment_name = exp.exp_name 50 | 51 | model = exp.get_model() 52 | if args.ckpt is None: 53 | file_name = os.path.join(exp.output_dir, args.experiment_name) 54 | ckpt_file = os.path.join(file_name, "best_ckpt.pth") 55 | else: 56 | ckpt_file = args.ckpt 57 | 58 | # load the model state dict 59 | ckpt = torch.load(ckpt_file, map_location="cpu") 60 | 61 | model.eval() 62 | if "model" in ckpt: 63 | ckpt = ckpt["model"] 64 | model.load_state_dict(ckpt) 65 | model.head.decode_in_inference = False 66 | 67 | logger.info("loading checkpoint done.") 68 | dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1]) 69 | 70 | mod = torch.jit.trace(model, dummy_input) 71 | mod.save(args.output_name) 72 | logger.info("generated torchscript model named {}".format(args.output_name)) 73 | 74 | 75 | if __name__ == "__main__": 76 | main() 77 | -------------------------------------------------------------------------------- /tools/process_trackingnet.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | """ 4 | This script is used to process the original trackingnet zip files 5 | TRAIN_0.zip ~ TRAIN_3.zip should be put under "datasets" dir 6 | """ 7 | if __name__ == "__main__": 8 | data_root = "datasets" 9 | set_ids = list(range(4)) 10 | """unzip""" 11 | for i in set_ids: 12 | ori_zip_file = os.path.join(data_root, "TRAIN_%d.zip"%i) 13 | unzip_dir = os.path.join(data_root, "TRAIN_%d"%i) 14 | if not os.path.exists(unzip_dir): 15 | os.makedirs(unzip_dir) 16 | print("unzipping %s"%ori_zip_file) 17 | os.system("unzip -qq %s -d %s" %(ori_zip_file, unzip_dir)) 18 | print("%s is done."%ori_zip_file) 19 | frame_root = os.path.join(unzip_dir, "frames") 20 | os.makedirs(frame_root) 21 | zip_dir = os.path.join(unzip_dir, "zips") 22 | sub_zips = os.listdir(zip_dir) 23 | for sub_zip in sub_zips: 24 | seq_name = sub_zip.replace(".zip", "") 25 | unzip_frame_dir = os.path.join(frame_root, seq_name) 26 | if not os.path.exists(unzip_frame_dir): 27 | os.makedirs(unzip_frame_dir) 28 | src_path = os.path.join(zip_dir, sub_zip) 29 | des_dir = unzip_frame_dir 30 | os.system("unzip -qq %s -d %s" %(src_path, des_dir)) 31 | print("%s is done." %sub_zip) 32 | os.system("rm -rf %s"%zip_dir) 33 | os.chdir(data_root) 34 | os.makedirs("TrackingNet") 35 | os.system("mv TRAIN_* TrackingNet") 36 | -------------------------------------------------------------------------------- /tools/test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | 5 | external_dir = os.path.join(os.path.dirname(__file__), "..", "external") 6 | if external_dir not in sys.path: 7 | sys.path.append(external_dir) 8 | 9 | 10 | from lib.test.evaluation import get_dataset 11 | from lib.test.evaluation.running import run_dataset 12 | from lib.test.evaluation.tracker import Tracker 13 | 14 | 15 | def run_tracker(tracker_name, tracker_param, run_id=None, dataset_name='otb', sequence=None, debug=0, threads=0, 16 | num_gpus=8): 17 | """Run tracker on sequence or dataset. 18 | args: 19 | tracker_name: Name of tracking method. 20 | tracker_param: Name of parameter file. 21 | run_id: The run id. 22 | dataset_name: Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot). 23 | sequence: Sequence number or name. 24 | debug: Debug level. 25 | threads: Number of threads. 26 | """ 27 | 28 | dataset = get_dataset(dataset_name) 29 | 30 | if sequence is not None: 31 | dataset = [dataset[sequence]] 32 | 33 | trackers = [Tracker(tracker_name, tracker_param, dataset_name, run_id)] 34 | 35 | run_dataset(dataset, trackers, debug, threads, num_gpus=num_gpus) 36 | 37 | 38 | def main(): 39 | parser = argparse.ArgumentParser(description='Run tracker on sequence or dataset.') 40 | parser.add_argument('tracker_name', type=str, help='Name of tracking method.') 41 | parser.add_argument('tracker_param', type=str, help='Name of config file.') 42 | parser.add_argument('--runid', type=int, default=None, help='The run id.') 43 | parser.add_argument('--dataset_name', type=str, default='otb', help='Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot).') 44 | parser.add_argument('--sequence', type=str, default=None, help='Sequence number or name.') 45 | parser.add_argument('--debug', type=int, default=0, help='Debug level.') 46 | parser.add_argument('--threads', type=int, default=0, help='Number of threads.') 47 | parser.add_argument('--num_gpus', type=int, default=8) 48 | 49 | args = parser.parse_args() 50 | 51 | try: 52 | seq_name = int(args.sequence) 53 | except: 54 | seq_name = args.sequence 55 | 56 | run_tracker(args.tracker_name, args.tracker_param, args.runid, args.dataset_name, seq_name, args.debug, 57 | args.threads, num_gpus=args.num_gpus) 58 | 59 | 60 | if __name__ == '__main__': 61 | main() 62 | -------------------------------------------------------------------------------- /unicorn/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | 4 | from .utils import configure_module 5 | 6 | configure_module() 7 | 8 | __version__ = "0.1.0" 9 | -------------------------------------------------------------------------------- /unicorn/core/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | from .launch import launch 6 | from .trainer import Trainer 7 | -------------------------------------------------------------------------------- /unicorn/data/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | from .data_augment import TrainTransform, ValTransform, TrainTransform_local, TrainTransform_omni, TrainTransform_Ins, TrainTransform_4tasks 6 | from .data_prefetcher import * 7 | from .dataloading import DataLoader, get_unicorn_datadir, worker_init_reset_seed 8 | # from .datasets import * 9 | from .samplers import InfiniteSampler, YoloBatchSampler 10 | -------------------------------------------------------------------------------- /unicorn/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | from .coco import COCODataset 6 | from .coco_classes import COCO_CLASSES 7 | from .datasets_wrapper import ConcatDataset, Dataset, MixConcatDataset 8 | from .mosaicdetection import MosaicDetection 9 | from .voc import VOCDetection 10 | # SOT 11 | from .coco_sot import COCOSOTDataset 12 | from .lasot import Lasot 13 | from .got10k import Got10k 14 | from .tracking_net import TrackingNet 15 | from .omni_data import OmniDataset 16 | # MOT 17 | from .mot import MOTDataset 18 | from .mot_omni import MOTOmniDataset 19 | from .bdd import BDDDataset 20 | from .bdd_omni import BDDOmniDataset 21 | # SOT-MOT 22 | from .mosaicdetection_uni import MosaicDetectionUni 23 | from .omni_data import OmniDatasetPlus 24 | # Instance Segmentation 25 | from .coco_inst import COCOInsDataset 26 | from .mosaicdetection import MosaicDetectionIns 27 | # MOTS 28 | from .coco_mots import COCOMOTSDataset 29 | from .mots_mot import MOTSMOTDataset 30 | from .bdd_omni_mots import BDDOmniMOTSDataset 31 | # VOS 32 | from .saliency import SaliencyDataset 33 | from .davis import DAVISDataset 34 | from .youtube_vos import YoutubeVOSDataset 35 | # SOT-MOT-VOS-MOTS 36 | from .mosaicdetection_uni import MosaicDetectionUni4tasks -------------------------------------------------------------------------------- /unicorn/data/datasets/coco_classes.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | COCO_CLASSES = ( 6 | "person", 7 | "bicycle", 8 | "car", 9 | "motorcycle", 10 | "airplane", 11 | "bus", 12 | "train", 13 | "truck", 14 | "boat", 15 | "traffic light", 16 | "fire hydrant", 17 | "stop sign", 18 | "parking meter", 19 | "bench", 20 | "bird", 21 | "cat", 22 | "dog", 23 | "horse", 24 | "sheep", 25 | "cow", 26 | "elephant", 27 | "bear", 28 | "zebra", 29 | "giraffe", 30 | "backpack", 31 | "umbrella", 32 | "handbag", 33 | "tie", 34 | "suitcase", 35 | "frisbee", 36 | "skis", 37 | "snowboard", 38 | "sports ball", 39 | "kite", 40 | "baseball bat", 41 | "baseball glove", 42 | "skateboard", 43 | "surfboard", 44 | "tennis racket", 45 | "bottle", 46 | "wine glass", 47 | "cup", 48 | "fork", 49 | "knife", 50 | "spoon", 51 | "bowl", 52 | "banana", 53 | "apple", 54 | "sandwich", 55 | "orange", 56 | "broccoli", 57 | "carrot", 58 | "hot dog", 59 | "pizza", 60 | "donut", 61 | "cake", 62 | "chair", 63 | "couch", 64 | "potted plant", 65 | "bed", 66 | "dining table", 67 | "toilet", 68 | "tv", 69 | "laptop", 70 | "mouse", 71 | "remote", 72 | "keyboard", 73 | "cell phone", 74 | "microwave", 75 | "oven", 76 | "toaster", 77 | "sink", 78 | "refrigerator", 79 | "book", 80 | "clock", 81 | "vase", 82 | "scissors", 83 | "teddy bear", 84 | "hair drier", 85 | "toothbrush", 86 | ) 87 | -------------------------------------------------------------------------------- /unicorn/data/datasets/voc_classes.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | # VOC_CLASSES = ( '__background__', # always index 0 6 | VOC_CLASSES = ( 7 | "aeroplane", 8 | "bicycle", 9 | "bird", 10 | "boat", 11 | "bottle", 12 | "bus", 13 | "car", 14 | "cat", 15 | "chair", 16 | "cow", 17 | "diningtable", 18 | "dog", 19 | "horse", 20 | "motorbike", 21 | "person", 22 | "pottedplant", 23 | "sheep", 24 | "sofa", 25 | "train", 26 | "tvmonitor", 27 | ) 28 | -------------------------------------------------------------------------------- /unicorn/data/samplers.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | import itertools 6 | from typing import Optional 7 | 8 | import torch 9 | import torch.distributed as dist 10 | from torch.utils.data.sampler import BatchSampler as torchBatchSampler 11 | from torch.utils.data.sampler import Sampler 12 | 13 | 14 | class YoloBatchSampler(torchBatchSampler): 15 | """ 16 | This batch sampler will generate mini-batches of (mosaic, index) tuples from another sampler. 17 | It works just like the :class:`torch.utils.data.sampler.BatchSampler`, 18 | but it will turn on/off the mosaic aug. 19 | """ 20 | 21 | def __init__(self, *args, mosaic=True, **kwargs): 22 | super().__init__(*args, **kwargs) 23 | self.mosaic = mosaic 24 | 25 | def __iter__(self): 26 | for batch in super().__iter__(): 27 | yield [(self.mosaic, idx) for idx in batch] 28 | 29 | 30 | class InfiniteSampler(Sampler): 31 | """ 32 | In training, we only care about the "infinite stream" of training data. 33 | So this sampler produces an infinite stream of indices and 34 | all workers cooperate to correctly shuffle the indices and sample different indices. 35 | The samplers in each worker effectively produces `indices[worker_id::num_workers]` 36 | where `indices` is an infinite stream of indices consisting of 37 | `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True) 38 | or `range(size) + range(size) + ...` (if shuffle is False) 39 | """ 40 | 41 | def __init__( 42 | self, 43 | size: int, 44 | shuffle: bool = True, 45 | seed: Optional[int] = 0, 46 | rank=0, 47 | world_size=1, 48 | ): 49 | """ 50 | Args: 51 | size (int): the total number of data of the underlying dataset to sample from 52 | shuffle (bool): whether to shuffle the indices or not 53 | seed (int): the initial seed of the shuffle. Must be the same 54 | across all workers. If None, will use a random seed shared 55 | among workers (require synchronization among all workers). 56 | """ 57 | self._size = size 58 | assert size > 0 59 | self._shuffle = shuffle 60 | self._seed = int(seed) 61 | 62 | if dist.is_available() and dist.is_initialized(): 63 | self._rank = dist.get_rank() 64 | self._world_size = dist.get_world_size() 65 | else: 66 | self._rank = rank 67 | self._world_size = world_size 68 | 69 | def __iter__(self): 70 | start = self._rank 71 | yield from itertools.islice( 72 | self._infinite_indices(), start, None, self._world_size 73 | ) 74 | 75 | def _infinite_indices(self): 76 | g = torch.Generator() 77 | g.manual_seed(self._seed) 78 | while True: 79 | if self._shuffle: 80 | yield from torch.randperm(self._size, generator=g) 81 | else: 82 | yield from torch.arange(self._size) 83 | 84 | def __len__(self): 85 | return self._size // self._world_size 86 | -------------------------------------------------------------------------------- /unicorn/evaluators/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | from .coco_evaluator import COCOEvaluator 6 | from .voc_evaluator import VOCEvaluator 7 | from .mot_evaluator import MOTEvaluator 8 | from .bdd_evaluator import BDDEvaluator 9 | from .coco_inst_evaluator import COCOInstEvaluator -------------------------------------------------------------------------------- /unicorn/exp/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | from .base_exp import BaseExp 6 | from .build import get_exp 7 | from .unicorn_det import ExpDet 8 | from .unicorn_det_mask import ExpDetMask 9 | from .unicorn_track import ExpTrack 10 | from .unicorn_track_mask import ExpTrackMask -------------------------------------------------------------------------------- /unicorn/exp/base_exp.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import ast 6 | import pprint 7 | from abc import ABCMeta, abstractmethod 8 | from typing import Dict 9 | from tabulate import tabulate 10 | 11 | import torch 12 | from torch.nn import Module 13 | 14 | from unicorn.utils import LRScheduler 15 | 16 | 17 | class BaseExp(metaclass=ABCMeta): 18 | """Basic class for any experiment.""" 19 | 20 | def __init__(self): 21 | self.seed = None 22 | self.output_dir = "./Unicorn_outputs" 23 | self.print_interval = 100 24 | self.eval_interval = 10 25 | 26 | @abstractmethod 27 | def get_model(self) -> Module: 28 | pass 29 | 30 | @abstractmethod 31 | def get_data_loader( 32 | self, batch_size: int, is_distributed: bool 33 | ) -> Dict[str, torch.utils.data.DataLoader]: 34 | pass 35 | 36 | @abstractmethod 37 | def get_optimizer(self, batch_size: int) -> torch.optim.Optimizer: 38 | pass 39 | 40 | @abstractmethod 41 | def get_lr_scheduler( 42 | self, lr: float, iters_per_epoch: int, **kwargs 43 | ) -> LRScheduler: 44 | pass 45 | 46 | @abstractmethod 47 | def get_evaluator(self): 48 | pass 49 | 50 | @abstractmethod 51 | def eval(self, model, evaluator, weights): 52 | pass 53 | 54 | def __repr__(self): 55 | table_header = ["keys", "values"] 56 | exp_table = [ 57 | (str(k), pprint.pformat(v)) 58 | for k, v in vars(self).items() 59 | if not k.startswith("_") 60 | ] 61 | return tabulate(exp_table, headers=table_header, tablefmt="fancy_grid") 62 | 63 | def merge(self, cfg_list): 64 | assert len(cfg_list) % 2 == 0 65 | for k, v in zip(cfg_list[0::2], cfg_list[1::2]): 66 | # only update value with same key 67 | if hasattr(self, k): 68 | src_value = getattr(self, k) 69 | src_type = type(src_value) 70 | if src_value is not None and src_type != type(v): 71 | try: 72 | v = src_type(v) 73 | except Exception: 74 | v = ast.literal_eval(v) 75 | setattr(self, k, v) 76 | -------------------------------------------------------------------------------- /unicorn/exp/build.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import importlib 6 | import os 7 | import sys 8 | 9 | 10 | def get_exp_by_file(exp_file): 11 | sys.path.append(os.path.dirname(exp_file)) 12 | current_exp = importlib.import_module(os.path.basename(exp_file).split(".")[0]) 13 | exp = current_exp.Exp() 14 | return exp 15 | 16 | 17 | def get_exp_by_name(exp_name): 18 | import yolox 19 | 20 | yolox_path = os.path.dirname(os.path.dirname(yolox.__file__)) 21 | filedict = { 22 | "yolox-s": "yolox_s.py", 23 | "yolox-m": "yolox_m.py", 24 | "yolox-l": "yolox_l.py", 25 | "yolox-x": "yolox_x.py", 26 | "yolox-tiny": "yolox_tiny.py", 27 | "yolox-nano": "nano.py", 28 | "yolov3": "yolov3.py", 29 | } 30 | filename = filedict[exp_name] 31 | exp_path = os.path.join(yolox_path, "exps", "default", filename) 32 | return get_exp_by_file(exp_path) 33 | 34 | 35 | def get_exp(exp_file, exp_name): 36 | """ 37 | get Exp object by file or name. If exp_file and exp_name 38 | are both provided, get Exp by exp_file. 39 | 40 | Args: 41 | exp_file (str): file path of experiment. 42 | exp_name (str): name of experiment. "yolo-s", 43 | """ 44 | assert ( 45 | exp_file is not None or exp_name is not None 46 | ), "plz provide exp file or exp name." 47 | if exp_file is not None: 48 | return get_exp_by_file(exp_file) 49 | else: 50 | return get_exp_by_name(exp_name) 51 | -------------------------------------------------------------------------------- /unicorn/layers/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | from .fast_coco_eval_api import COCOeval_opt 6 | -------------------------------------------------------------------------------- /unicorn/layers/csrc/vision.cpp: -------------------------------------------------------------------------------- 1 | #include "cocoeval/cocoeval.h" 2 | 3 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 4 | m.def("COCOevalAccumulate", &COCOeval::Accumulate, "COCOeval::Accumulate"); 5 | m.def( 6 | "COCOevalEvaluateImages", 7 | &COCOeval::EvaluateImages, 8 | "COCOeval::EvaluateImages"); 9 | pybind11::class_(m, "InstanceAnnotation") 10 | .def(pybind11::init()); 11 | pybind11::class_(m, "ImageEvaluation") 12 | .def(pybind11::init<>()); 13 | } 14 | -------------------------------------------------------------------------------- /unicorn/models/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | # Object Detection 6 | from .losses import IOUloss 7 | from .backbone.darknet import CSPDarknet, Darknet 8 | from .backbone.yolo_pafpn_new import YOLOPAFPNNEW 9 | from .yolo_head_det import YOLOXHeadDet 10 | from .yolox import YOLOX 11 | # Instance Segmentation 12 | from .yolo_head_det_mask import YOLOXHeadDetMask 13 | from .yolox import YOLOXMask 14 | # Unified Tracking and Segmentation 15 | from .unicorn_head import UnicornHead 16 | from .unicorn_head_mask import UnicornHeadMask 17 | from .unicorn import Unicorn, UnicornActor 18 | -------------------------------------------------------------------------------- /unicorn/models/condinst/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /unicorn/models/condinst/comm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def aligned_bilinear(tensor, factor): 6 | assert tensor.dim() == 4 7 | assert factor >= 1 8 | assert int(factor) == factor 9 | 10 | if factor == 1: 11 | return tensor 12 | 13 | h, w = tensor.size()[2:] 14 | tensor = F.pad(tensor, pad=(0, 1, 0, 1), mode="replicate") 15 | oh = factor * h + 1 16 | ow = factor * w + 1 17 | tensor = F.interpolate( 18 | tensor, size=(oh, ow), 19 | mode='bilinear', 20 | align_corners=True 21 | ) 22 | tensor = F.pad( 23 | tensor, pad=(factor // 2, 0, factor // 2, 0), 24 | mode="replicate" 25 | ) 26 | 27 | return tensor[:, :, :oh - 1, :ow - 1] 28 | 29 | 30 | def compute_locations(h, w, stride, device): 31 | shifts_x = torch.arange( 32 | 0, w * stride, step=stride, 33 | dtype=torch.float32, device=device 34 | ) 35 | shifts_y = torch.arange( 36 | 0, h * stride, step=stride, 37 | dtype=torch.float32, device=device 38 | ) 39 | shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x) 40 | shift_x = shift_x.reshape(-1) 41 | shift_y = shift_y.reshape(-1) 42 | locations = torch.stack((shift_x, shift_y), dim=1) + stride // 2 # (hw, 2) 43 | return locations 44 | 45 | 46 | def compute_ious(pred, target): 47 | """ 48 | Args: 49 | pred: Nx4 predicted bounding boxes 50 | target: Nx4 target bounding boxes 51 | Both are in the form of FCOS prediction (l, t, r, b) 52 | """ 53 | pred_left = pred[:, 0] 54 | pred_top = pred[:, 1] 55 | pred_right = pred[:, 2] 56 | pred_bottom = pred[:, 3] 57 | 58 | target_left = target[:, 0] 59 | target_top = target[:, 1] 60 | target_right = target[:, 2] 61 | target_bottom = target[:, 3] 62 | 63 | target_aera = (target_left + target_right) * \ 64 | (target_top + target_bottom) 65 | pred_aera = (pred_left + pred_right) * \ 66 | (pred_top + pred_bottom) 67 | 68 | w_intersect = torch.min(pred_left, target_left) + \ 69 | torch.min(pred_right, target_right) 70 | h_intersect = torch.min(pred_bottom, target_bottom) + \ 71 | torch.min(pred_top, target_top) 72 | 73 | g_w_intersect = torch.max(pred_left, target_left) + \ 74 | torch.max(pred_right, target_right) 75 | g_h_intersect = torch.max(pred_bottom, target_bottom) + \ 76 | torch.max(pred_top, target_top) 77 | ac_uion = g_w_intersect * g_h_intersect 78 | 79 | area_intersect = w_intersect * h_intersect 80 | area_union = target_aera + pred_aera - area_intersect 81 | 82 | ious = (area_intersect + 1.0) / (area_union + 1.0) 83 | gious = ious - (ac_uion - area_union) / ac_uion 84 | 85 | return ious, gious 86 | -------------------------------------------------------------------------------- /unicorn/models/condinst/config/__init__.py: -------------------------------------------------------------------------------- 1 | from .config import get_cfg 2 | 3 | __all__ = [ 4 | "get_cfg", 5 | ] 6 | -------------------------------------------------------------------------------- /unicorn/models/condinst/config/config.py: -------------------------------------------------------------------------------- 1 | # from detectron2.config import CfgNode 2 | import copy 3 | 4 | def get_cfg(): 5 | """ 6 | Get a copy of the default config. 7 | 8 | Returns: 9 | a detectron2 CfgNode instance. 10 | """ 11 | from .defaults import _C 12 | 13 | return copy.deepcopy(_C) 14 | -------------------------------------------------------------------------------- /unicorn/models/losses.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- encoding: utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import torch 6 | import torch.nn as nn 7 | 8 | 9 | class IOUloss(nn.Module): 10 | def __init__(self, reduction="none", loss_type="iou"): 11 | super(IOUloss, self).__init__() 12 | self.reduction = reduction 13 | self.loss_type = loss_type 14 | 15 | def forward(self, pred, target): 16 | assert pred.shape[0] == target.shape[0] 17 | 18 | pred = pred.view(-1, 4) 19 | target = target.view(-1, 4) 20 | tl = torch.max( 21 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2) 22 | ) 23 | br = torch.min( 24 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2) 25 | ) 26 | 27 | area_p = torch.prod(pred[:, 2:], 1) 28 | area_g = torch.prod(target[:, 2:], 1) 29 | 30 | en = (tl < br).type(tl.type()).prod(dim=1) 31 | area_i = torch.prod(br - tl, 1) * en 32 | area_u = area_p + area_g - area_i 33 | iou = (area_i) / (area_u + 1e-16) 34 | 35 | if self.loss_type == "iou": 36 | loss = 1 - iou ** 2 37 | elif self.loss_type == "giou": 38 | c_tl = torch.min( 39 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2) 40 | ) 41 | c_br = torch.max( 42 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2) 43 | ) 44 | area_c = torch.prod(c_br - c_tl, 1) 45 | giou = iou - (area_c - area_u) / area_c.clamp(1e-16) 46 | loss = 1 - giou.clamp(min=-1.0, max=1.0) 47 | 48 | if self.reduction == "mean": 49 | loss = loss.mean() 50 | elif self.reduction == "sum": 51 | loss = loss.sum() 52 | 53 | return loss 54 | -------------------------------------------------------------------------------- /unicorn/models/ops/functions/__init__.py: -------------------------------------------------------------------------------- 1 | # ------------------------------------------------------------------------------------------------ 2 | # Deformable DETR 3 | # Copyright (c) 2020 SenseTime. All Rights Reserved. 4 | # Licensed under the Apache License, Version 2.0 [see LICENSE for details] 5 | # ------------------------------------------------------------------------------------------------ 6 | # Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 7 | # ------------------------------------------------------------------------------------------------ 8 | 9 | from .ms_deform_attn_func import MSDeformAttnFunction 10 | 11 | -------------------------------------------------------------------------------- /unicorn/models/ops/make.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # ------------------------------------------------------------------------------------------------ 3 | # Deformable DETR 4 | # Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | # Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | # ------------------------------------------------------------------------------------------------ 7 | # Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | # ------------------------------------------------------------------------------------------------ 9 | 10 | python3 setup.py build install --user 11 | -------------------------------------------------------------------------------- /unicorn/models/ops/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # ------------------------------------------------------------------------------------------------ 2 | # Deformable DETR 3 | # Copyright (c) 2020 SenseTime. All Rights Reserved. 4 | # Licensed under the Apache License, Version 2.0 [see LICENSE for details] 5 | # ------------------------------------------------------------------------------------------------ 6 | # Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 7 | # ------------------------------------------------------------------------------------------------ 8 | 9 | from .ms_deform_attn import MSDeformAttn 10 | -------------------------------------------------------------------------------- /unicorn/models/ops/setup.py: -------------------------------------------------------------------------------- 1 | # ------------------------------------------------------------------------------------------------ 2 | # Deformable DETR 3 | # Copyright (c) 2020 SenseTime. All Rights Reserved. 4 | # Licensed under the Apache License, Version 2.0 [see LICENSE for details] 5 | # ------------------------------------------------------------------------------------------------ 6 | # Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 7 | # ------------------------------------------------------------------------------------------------ 8 | 9 | import os 10 | import glob 11 | 12 | import torch 13 | 14 | from torch.utils.cpp_extension import CUDA_HOME 15 | from torch.utils.cpp_extension import CppExtension 16 | from torch.utils.cpp_extension import CUDAExtension 17 | 18 | from setuptools import find_packages 19 | from setuptools import setup 20 | 21 | requirements = ["torch", "torchvision"] 22 | 23 | def get_extensions(): 24 | this_dir = os.path.dirname(os.path.abspath(__file__)) 25 | extensions_dir = os.path.join(this_dir, "src") 26 | 27 | main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) 28 | source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp")) 29 | source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu")) 30 | 31 | sources = main_file + source_cpu 32 | extension = CppExtension 33 | extra_compile_args = {"cxx": []} 34 | define_macros = [] 35 | 36 | if torch.cuda.is_available() and CUDA_HOME is not None: 37 | extension = CUDAExtension 38 | sources += source_cuda 39 | define_macros += [("WITH_CUDA", None)] 40 | extra_compile_args["nvcc"] = [ 41 | "-DCUDA_HAS_FP16=1", 42 | "-D__CUDA_NO_HALF_OPERATORS__", 43 | "-D__CUDA_NO_HALF_CONVERSIONS__", 44 | "-D__CUDA_NO_HALF2_OPERATORS__", 45 | ] 46 | else: 47 | raise NotImplementedError('Cuda is not availabel') 48 | 49 | sources = [os.path.join(extensions_dir, s) for s in sources] 50 | include_dirs = [extensions_dir] 51 | ext_modules = [ 52 | extension( 53 | "MultiScaleDeformableAttention", 54 | sources, 55 | include_dirs=include_dirs, 56 | define_macros=define_macros, 57 | extra_compile_args=extra_compile_args, 58 | ) 59 | ] 60 | return ext_modules 61 | 62 | setup( 63 | name="MultiScaleDeformableAttention", 64 | version="1.0", 65 | author="Weijie Su", 66 | url="https://github.com/fundamentalvision/Deformable-DETR", 67 | description="PyTorch Wrapper for CUDA Functions of Multi-Scale Deformable Attention", 68 | packages=find_packages(exclude=("configs", "tests",)), 69 | ext_modules=get_extensions(), 70 | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, 71 | ) 72 | -------------------------------------------------------------------------------- /unicorn/models/ops/src/cpu/ms_deform_attn_cpu.cpp: -------------------------------------------------------------------------------- 1 | /*! 2 | ************************************************************************************************** 3 | * Deformable DETR 4 | * Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | * Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | ************************************************************************************************** 7 | * Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | ************************************************************************************************** 9 | */ 10 | 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | 17 | at::Tensor 18 | ms_deform_attn_cpu_forward( 19 | const at::Tensor &value, 20 | const at::Tensor &spatial_shapes, 21 | const at::Tensor &level_start_index, 22 | const at::Tensor &sampling_loc, 23 | const at::Tensor &attn_weight, 24 | const int im2col_step) 25 | { 26 | AT_ERROR("Not implement on cpu"); 27 | } 28 | 29 | std::vector 30 | ms_deform_attn_cpu_backward( 31 | const at::Tensor &value, 32 | const at::Tensor &spatial_shapes, 33 | const at::Tensor &level_start_index, 34 | const at::Tensor &sampling_loc, 35 | const at::Tensor &attn_weight, 36 | const at::Tensor &grad_output, 37 | const int im2col_step) 38 | { 39 | AT_ERROR("Not implement on cpu"); 40 | } 41 | 42 | -------------------------------------------------------------------------------- /unicorn/models/ops/src/cpu/ms_deform_attn_cpu.h: -------------------------------------------------------------------------------- 1 | /*! 2 | ************************************************************************************************** 3 | * Deformable DETR 4 | * Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | * Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | ************************************************************************************************** 7 | * Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | ************************************************************************************************** 9 | */ 10 | 11 | #pragma once 12 | #include 13 | 14 | at::Tensor 15 | ms_deform_attn_cpu_forward( 16 | const at::Tensor &value, 17 | const at::Tensor &spatial_shapes, 18 | const at::Tensor &level_start_index, 19 | const at::Tensor &sampling_loc, 20 | const at::Tensor &attn_weight, 21 | const int im2col_step); 22 | 23 | std::vector 24 | ms_deform_attn_cpu_backward( 25 | const at::Tensor &value, 26 | const at::Tensor &spatial_shapes, 27 | const at::Tensor &level_start_index, 28 | const at::Tensor &sampling_loc, 29 | const at::Tensor &attn_weight, 30 | const at::Tensor &grad_output, 31 | const int im2col_step); 32 | 33 | 34 | -------------------------------------------------------------------------------- /unicorn/models/ops/src/cuda/ms_deform_attn_cuda.h: -------------------------------------------------------------------------------- 1 | /*! 2 | ************************************************************************************************** 3 | * Deformable DETR 4 | * Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | * Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | ************************************************************************************************** 7 | * Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | ************************************************************************************************** 9 | */ 10 | 11 | #pragma once 12 | #include 13 | 14 | at::Tensor ms_deform_attn_cuda_forward( 15 | const at::Tensor &value, 16 | const at::Tensor &spatial_shapes, 17 | const at::Tensor &level_start_index, 18 | const at::Tensor &sampling_loc, 19 | const at::Tensor &attn_weight, 20 | const int im2col_step); 21 | 22 | std::vector ms_deform_attn_cuda_backward( 23 | const at::Tensor &value, 24 | const at::Tensor &spatial_shapes, 25 | const at::Tensor &level_start_index, 26 | const at::Tensor &sampling_loc, 27 | const at::Tensor &attn_weight, 28 | const at::Tensor &grad_output, 29 | const int im2col_step); 30 | 31 | -------------------------------------------------------------------------------- /unicorn/models/ops/src/ms_deform_attn.h: -------------------------------------------------------------------------------- 1 | /*! 2 | ************************************************************************************************** 3 | * Deformable DETR 4 | * Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | * Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | ************************************************************************************************** 7 | * Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | ************************************************************************************************** 9 | */ 10 | 11 | #pragma once 12 | 13 | #include "cpu/ms_deform_attn_cpu.h" 14 | 15 | #ifdef WITH_CUDA 16 | #include "cuda/ms_deform_attn_cuda.h" 17 | #endif 18 | 19 | 20 | at::Tensor 21 | ms_deform_attn_forward( 22 | const at::Tensor &value, 23 | const at::Tensor &spatial_shapes, 24 | const at::Tensor &level_start_index, 25 | const at::Tensor &sampling_loc, 26 | const at::Tensor &attn_weight, 27 | const int im2col_step) 28 | { 29 | if (value.type().is_cuda()) 30 | { 31 | #ifdef WITH_CUDA 32 | return ms_deform_attn_cuda_forward( 33 | value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); 34 | #else 35 | AT_ERROR("Not compiled with GPU support"); 36 | #endif 37 | } 38 | AT_ERROR("Not implemented on the CPU"); 39 | } 40 | 41 | std::vector 42 | ms_deform_attn_backward( 43 | const at::Tensor &value, 44 | const at::Tensor &spatial_shapes, 45 | const at::Tensor &level_start_index, 46 | const at::Tensor &sampling_loc, 47 | const at::Tensor &attn_weight, 48 | const at::Tensor &grad_output, 49 | const int im2col_step) 50 | { 51 | if (value.type().is_cuda()) 52 | { 53 | #ifdef WITH_CUDA 54 | return ms_deform_attn_cuda_backward( 55 | value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step); 56 | #else 57 | AT_ERROR("Not compiled with GPU support"); 58 | #endif 59 | } 60 | AT_ERROR("Not implemented on the CPU"); 61 | } 62 | 63 | -------------------------------------------------------------------------------- /unicorn/models/ops/src/vision.cpp: -------------------------------------------------------------------------------- 1 | /*! 2 | ************************************************************************************************** 3 | * Deformable DETR 4 | * Copyright (c) 2020 SenseTime. All Rights Reserved. 5 | * Licensed under the Apache License, Version 2.0 [see LICENSE for details] 6 | ************************************************************************************************** 7 | * Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 8 | ************************************************************************************************** 9 | */ 10 | 11 | #include "ms_deform_attn.h" 12 | 13 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 14 | m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward"); 15 | m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward"); 16 | } 17 | -------------------------------------------------------------------------------- /unicorn/models/position_encoding.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various positional encodings for the transformer. 3 | """ 4 | import math 5 | import torch 6 | from torch import nn 7 | import torch.nn.functional as F 8 | 9 | 10 | class PositionEmbeddingLearned(nn.Module): 11 | """ 12 | Absolute pos embedding, learned. (allow users to specify the size) 13 | """ 14 | def __init__(self, num_pos_feats=256, sz=20): 15 | super().__init__() 16 | self.sz = sz 17 | self.row_embed = nn.Embedding(sz, num_pos_feats) 18 | self.col_embed = nn.Embedding(sz, num_pos_feats) 19 | self.reset_parameters() 20 | 21 | def reset_parameters(self): 22 | nn.init.uniform_(self.row_embed.weight) 23 | nn.init.uniform_(self.col_embed.weight) 24 | 25 | def forward(self, bs, dh, dw): 26 | """bs: batch size, dh: destination h, dw: destination w""" 27 | h, w = self.sz, self.sz 28 | i = torch.arange(w, device=self.col_embed.weight.device) 29 | j = torch.arange(h, device=self.row_embed.weight.device) 30 | x_emb = self.col_embed(i) 31 | y_emb = self.row_embed(j) 32 | pos = torch.cat([ 33 | x_emb.unsqueeze(0).repeat(h, 1, 1), 34 | y_emb.unsqueeze(1).repeat(1, w, 1), 35 | ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(bs, 1, 1, 1) # (H,W,C) --> (C,H,W) --> (1,C,H,W) --> (B,C,H,W) 36 | return F.interpolate(pos, (dh, dw), mode="bilinear", align_corners=False) 37 | 38 | 39 | def build_position_encoding(hidden_dim=256, sz=40): 40 | N_steps = hidden_dim // 2 41 | position_embedding = PositionEmbeddingLearned(N_steps, sz) 42 | return position_embedding 43 | -------------------------------------------------------------------------------- /unicorn/models/yolox.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- encoding: utf-8 -*- 3 | # ------------------------------------------------------------------------ 4 | # Unicorn 5 | # Copyright (c) 2022 ByteDance. All Rights Reserved. 6 | # Licensed under the MIT License [see LICENSE for details] 7 | # ------------------------------------------------------------------------ 8 | # Modified from YOLOX (https://github.com/Megvii-BaseDetection/YOLOX) 9 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 10 | # ---------------------------------------------------------------------- 11 | 12 | import torch.nn as nn 13 | 14 | # import time 15 | 16 | class YOLOX(nn.Module): 17 | """ 18 | YOLOX model module. The module list is defined by create_yolov3_modules function. 19 | The network returns loss values from three YOLO layers during training 20 | and detection results during test. 21 | """ 22 | 23 | def __init__(self, backbone=None, head=None): 24 | super().__init__() 25 | 26 | self.backbone = backbone 27 | self.head = head 28 | 29 | def forward(self, x, targets=None): 30 | # fpn output content features of [dark3, dark4, dark5] 31 | # s = time.time() 32 | fpn_outs = self.backbone(x) 33 | # e_b = time.time() 34 | if self.training: 35 | assert targets is not None 36 | loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.head( 37 | fpn_outs, targets, x 38 | ) 39 | outputs = { 40 | "total_loss": loss, 41 | "iou_loss": iou_loss, 42 | "l1_loss": l1_loss, 43 | "conf_loss": conf_loss, 44 | "cls_loss": cls_loss, 45 | "num_fg": num_fg, 46 | } 47 | else: 48 | outputs = self.head(fpn_outs) 49 | # e_d = time.time() 50 | # print("backbone time: %.4f, head time: %.4f " % (e_b-s, e_d-e_b)) 51 | return outputs 52 | 53 | class YOLOXMask(nn.Module): 54 | """ 55 | YOLOX model module. The module list is defined by create_yolov3_modules function. 56 | The network returns loss values from three YOLO layers during training 57 | and detection results during test. 58 | """ 59 | 60 | def __init__(self, backbone=None, head=None): 61 | super().__init__() 62 | 63 | self.backbone = backbone 64 | self.head = head 65 | 66 | def forward(self, x, targets=None, masks=None): 67 | # fpn output content features of [dark3, dark4, dark5] 68 | # s = time.time() 69 | fpn_outs = self.backbone(x) 70 | # e_b = time.time() 71 | if self.training: 72 | assert targets is not None 73 | assert masks is not None 74 | loss_dict = self.head(fpn_outs, targets, x, masks) 75 | return loss_dict 76 | else: 77 | outputs = self.head(fpn_outs) 78 | # e_d = time.time() 79 | # print("backbone time: %.4f, head time: %.4f " % (e_b-s, e_d-e_b)) 80 | return outputs 81 | -------------------------------------------------------------------------------- /unicorn/tracker/basetrack.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from collections import OrderedDict 3 | 4 | 5 | class TrackState(object): 6 | New = 0 7 | Tracked = 1 8 | Lost = 2 9 | Removed = 3 10 | 11 | 12 | class BaseTrack(object): 13 | _count = 0 14 | 15 | track_id = 0 16 | is_activated = False 17 | state = TrackState.New 18 | 19 | history = OrderedDict() 20 | features = [] 21 | curr_feature = None 22 | score = 0 23 | start_frame = 0 24 | frame_id = 0 25 | time_since_update = 0 26 | 27 | # multi-camera 28 | location = (np.inf, np.inf) 29 | 30 | @property 31 | def end_frame(self): 32 | return self.frame_id 33 | 34 | @staticmethod 35 | def next_id(): 36 | BaseTrack._count += 1 37 | return BaseTrack._count 38 | 39 | def activate(self, *args): 40 | raise NotImplementedError 41 | 42 | def predict(self): 43 | raise NotImplementedError 44 | 45 | def update(self, *args, **kwargs): 46 | raise NotImplementedError 47 | 48 | def mark_lost(self): 49 | self.state = TrackState.Lost 50 | 51 | def mark_removed(self): 52 | self.state = TrackState.Removed 53 | 54 | @staticmethod 55 | def clean_id(): 56 | BaseTrack._count = 0 -------------------------------------------------------------------------------- /unicorn/utils/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | from .allreduce_norm import * 6 | from .boxes import * 7 | from .checkpoint import load_ckpt, save_checkpoint 8 | from .demo_utils import * 9 | from .dist import * 10 | from .ema import * 11 | from .logger import setup_logger 12 | from .lr_scheduler import LRScheduler 13 | from .metric import * 14 | from .model_utils import * 15 | from .setup_env import * 16 | from .visualize import * 17 | -------------------------------------------------------------------------------- /unicorn/utils/allreduce_norm.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import pickle 6 | from collections import OrderedDict 7 | 8 | import torch 9 | from torch import distributed as dist 10 | from torch import nn 11 | 12 | from .dist import _get_global_gloo_group, get_world_size 13 | 14 | ASYNC_NORM = ( 15 | nn.BatchNorm1d, 16 | nn.BatchNorm2d, 17 | nn.BatchNorm3d, 18 | nn.InstanceNorm1d, 19 | nn.InstanceNorm2d, 20 | nn.InstanceNorm3d, 21 | ) 22 | 23 | __all__ = [ 24 | "get_async_norm_states", 25 | "pyobj2tensor", 26 | "tensor2pyobj", 27 | "all_reduce", 28 | "all_reduce_norm", 29 | ] 30 | 31 | 32 | def get_async_norm_states(module): 33 | async_norm_states = OrderedDict() 34 | for name, child in module.named_modules(): 35 | if isinstance(child, ASYNC_NORM): 36 | for k, v in child.state_dict().items(): 37 | async_norm_states[".".join([name, k])] = v 38 | return async_norm_states 39 | 40 | 41 | def pyobj2tensor(pyobj, device="cuda"): 42 | """serialize picklable python object to tensor""" 43 | storage = torch.ByteStorage.from_buffer(pickle.dumps(pyobj)) 44 | return torch.ByteTensor(storage).to(device=device) 45 | 46 | 47 | def tensor2pyobj(tensor): 48 | """deserialize tensor to picklable python object""" 49 | return pickle.loads(tensor.cpu().numpy().tobytes()) 50 | 51 | 52 | def _get_reduce_op(op_name): 53 | return { 54 | "sum": dist.ReduceOp.SUM, 55 | "mean": dist.ReduceOp.SUM, 56 | }[op_name.lower()] 57 | 58 | 59 | def all_reduce(py_dict, op="sum", group=None): 60 | """ 61 | Apply all reduce function for python dict object. 62 | NOTE: make sure that every py_dict has the same keys and values are in the same shape. 63 | 64 | Args: 65 | py_dict (dict): dict to apply all reduce op. 66 | op (str): operator, could be "sum" or "mean". 67 | """ 68 | world_size = get_world_size() 69 | if world_size == 1: 70 | return py_dict 71 | if group is None: 72 | group = _get_global_gloo_group() 73 | if dist.get_world_size(group) == 1: 74 | return py_dict 75 | 76 | # all reduce logic across different devices. 77 | py_key = list(py_dict.keys()) 78 | py_key_tensor = pyobj2tensor(py_key) 79 | dist.broadcast(py_key_tensor, src=0) 80 | py_key = tensor2pyobj(py_key_tensor) 81 | 82 | tensor_shapes = [py_dict[k].shape for k in py_key] 83 | tensor_numels = [py_dict[k].numel() for k in py_key] 84 | 85 | flatten_tensor = torch.cat([py_dict[k].flatten() for k in py_key]) 86 | dist.all_reduce(flatten_tensor, op=_get_reduce_op(op)) 87 | if op == "mean": 88 | flatten_tensor /= world_size 89 | 90 | split_tensors = [ 91 | x.reshape(shape) 92 | for x, shape in zip(torch.split(flatten_tensor, tensor_numels), tensor_shapes) 93 | ] 94 | return OrderedDict({k: v for k, v in zip(py_key, split_tensors)}) 95 | 96 | 97 | def all_reduce_norm(module): 98 | """ 99 | All reduce norm statistics in different devices. 100 | """ 101 | states = get_async_norm_states(module) 102 | states = all_reduce(states, op="mean") 103 | module.load_state_dict(states, strict=False) 104 | -------------------------------------------------------------------------------- /unicorn/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | import os 5 | import shutil 6 | from loguru import logger 7 | 8 | import torch 9 | 10 | 11 | def load_ckpt(model, ckpt): 12 | model_state_dict = model.state_dict() 13 | load_dict = {} 14 | for key_model, v in model_state_dict.items(): 15 | if key_model not in ckpt: 16 | logger.warning( 17 | "{} is not in the ckpt. Please double check and see if this is desired.".format( 18 | key_model 19 | ) 20 | ) 21 | continue 22 | v_ckpt = ckpt[key_model] 23 | if v.shape != v_ckpt.shape: 24 | logger.warning( 25 | "Shape of {} in checkpoint is {}, while shape of {} in model is {}.".format( 26 | key_model, v_ckpt.shape, key_model, v.shape 27 | ) 28 | ) 29 | continue 30 | load_dict[key_model] = v_ckpt 31 | 32 | model.load_state_dict(load_dict, strict=False) 33 | return model 34 | 35 | 36 | def save_checkpoint(state, is_best, save_dir, model_name=""): 37 | if not os.path.exists(save_dir): 38 | os.makedirs(save_dir) 39 | filename = os.path.join(save_dir, model_name + "_ckpt.pth") 40 | torch.save(state, filename) 41 | if is_best: 42 | best_filename = os.path.join(save_dir, "best_ckpt.pth") 43 | shutil.copyfile(filename, best_filename) 44 | -------------------------------------------------------------------------------- /unicorn/utils/ema.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | import math 5 | from copy import deepcopy 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | __all__ = ["ModelEMA", "is_parallel"] 11 | 12 | 13 | def is_parallel(model): 14 | """check if model is in parallel mode.""" 15 | parallel_type = ( 16 | nn.parallel.DataParallel, 17 | nn.parallel.DistributedDataParallel, 18 | ) 19 | return isinstance(model, parallel_type) 20 | 21 | 22 | class ModelEMA: 23 | """ 24 | Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 25 | Keep a moving average of everything in the model state_dict (parameters and buffers). 26 | This is intended to allow functionality like 27 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 28 | A smoothed version of the weights is necessary for some training schemes to perform well. 29 | This class is sensitive where it is initialized in the sequence of model init, 30 | GPU assignment and distributed training wrappers. 31 | """ 32 | 33 | def __init__(self, model_wo_ddp, decay=0.9999, updates=0, is_distributed=True): 34 | """ 35 | Args: 36 | model (nn.Module): model to apply EMA. 37 | decay (float): ema decay reate. 38 | updates (int): counter of EMA updates. 39 | """ 40 | # Create EMA(FP32) 41 | self.is_distributed = is_distributed 42 | self.ema = deepcopy(model_wo_ddp).eval() 43 | # self.ema = deepcopy(model.module if self.is_distributed else model).eval() # deepcopy(model.module) may trigger some errors 44 | self.updates = updates 45 | # decay exponential ramp (to help early epochs) 46 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) 47 | for p in self.ema.parameters(): 48 | p.requires_grad_(False) 49 | 50 | def update(self, model): 51 | # Update EMA parameters 52 | with torch.no_grad(): 53 | self.updates += 1 54 | d = self.decay(self.updates) 55 | 56 | msd = ( 57 | model.module.state_dict() if self.is_distributed else model.state_dict() 58 | ) # model state_dict 59 | 60 | for k, v in self.ema.state_dict().items(): 61 | if v.dtype.is_floating_point: 62 | v *= d 63 | v += (1.0 - d) * msd[k].detach() 64 | -------------------------------------------------------------------------------- /unicorn/utils/logger.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import inspect 6 | import os 7 | import sys 8 | from loguru import logger 9 | 10 | 11 | def get_caller_name(depth=0): 12 | """ 13 | Args: 14 | depth (int): Depth of caller conext, use 0 for caller depth. Default value: 0. 15 | 16 | Returns: 17 | str: module name of the caller 18 | """ 19 | # the following logic is a little bit faster than inspect.stack() logic 20 | frame = inspect.currentframe().f_back 21 | for _ in range(depth): 22 | frame = frame.f_back 23 | 24 | return frame.f_globals["__name__"] 25 | 26 | 27 | class StreamToLoguru: 28 | """ 29 | stream object that redirects writes to a logger instance. 30 | """ 31 | 32 | def __init__(self, level="INFO", caller_names=("apex", "pycocotools")): 33 | """ 34 | Args: 35 | level(str): log level string of loguru. Default value: "INFO". 36 | caller_names(tuple): caller names of redirected module. 37 | Default value: (apex, pycocotools). 38 | """ 39 | self.level = level 40 | self.linebuf = "" 41 | self.caller_names = caller_names 42 | 43 | def write(self, buf): 44 | full_name = get_caller_name(depth=1) 45 | module_name = full_name.rsplit(".", maxsplit=-1)[0] 46 | if module_name in self.caller_names: 47 | for line in buf.rstrip().splitlines(): 48 | # use caller level log 49 | logger.opt(depth=2).log(self.level, line.rstrip()) 50 | else: 51 | sys.__stdout__.write(buf) 52 | 53 | def flush(self): 54 | pass 55 | 56 | 57 | def redirect_sys_output(log_level="INFO"): 58 | redirect_logger = StreamToLoguru(log_level) 59 | sys.stderr = redirect_logger 60 | sys.stdout = redirect_logger 61 | 62 | 63 | def setup_logger(save_dir, distributed_rank=0, filename="log.txt", mode="a"): 64 | """setup logger for training and testing. 65 | Args: 66 | save_dir(str): location to save log file 67 | distributed_rank(int): device rank when multi-gpu environment 68 | filename (string): log save name. 69 | mode(str): log file write mode, `append` or `override`. default is `a`. 70 | 71 | Return: 72 | logger instance. 73 | """ 74 | loguru_format = ( 75 | "{time:YYYY-MM-DD HH:mm:ss} | " 76 | "{level: <8} | " 77 | "{name}:{line} - {message}" 78 | ) 79 | 80 | logger.remove() 81 | save_file = os.path.join(save_dir, filename) 82 | if mode == "o" and os.path.exists(save_file): 83 | os.remove(save_file) 84 | # only keep logger in rank0 process 85 | if distributed_rank == 0: 86 | logger.add( 87 | sys.stderr, 88 | format=loguru_format, 89 | level="INFO", 90 | enqueue=True, 91 | ) 92 | logger.add(save_file) 93 | 94 | # redirect stdout/stderr to loguru 95 | redirect_sys_output("INFO") 96 | -------------------------------------------------------------------------------- /unicorn/utils/merge.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as NNF 3 | 4 | 5 | def merge_backbone_output(inp_list): 6 | """ merge outputs from backbone 7 | feat: HWxBxC, pos: HWxBxC, mask: BxHW 8 | """ 9 | nf = len(inp_list) 10 | seq_dict = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0), 11 | "mask": torch.cat([x["mask"] for x in inp_list], dim=1), 12 | "pos": torch.cat([x["pos"] for x in inp_list], dim=0), 13 | "h": inp_list[0]["h"], "w": inp_list[0]["w"], "nf": nf} 14 | return seq_dict 15 | 16 | 17 | def convert_to_onehot(x, dim): 18 | """convert to one-hot""" 19 | out_lbs = torch.zeros_like(x, device=x.device) 20 | pos_index = torch.argmax(x, dim=dim, keepdim=True) 21 | out_lbs.scatter_(dim, pos_index, 1) 22 | return out_lbs 23 | 24 | def adjust_labels_sz(inp_lbs, dh, dw): 25 | """ inp_lbs: PyTorch tensor (F, K, H, W) --> (F, K, H*s, W*s) 26 | dh: destination h, dw: destination w""" 27 | assert len(inp_lbs.size()) == 4 28 | """interpolation""" 29 | x = NNF.interpolate(inp_lbs, size=(dh, dw), mode="bilinear", align_corners=False) 30 | """convert to one-hot""" 31 | out_lbs = convert_to_onehot(x, dim=1) 32 | return out_lbs 33 | -------------------------------------------------------------------------------- /unicorn/utils/setup_env.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import os 6 | import subprocess 7 | from loguru import logger 8 | 9 | import cv2 10 | 11 | from .dist import get_world_size, is_main_process 12 | 13 | __all__ = ["configure_nccl", "configure_module", "configure_omp"] 14 | 15 | 16 | def configure_nccl(): 17 | """Configure multi-machine environment variables of NCCL.""" 18 | os.environ["NCCL_LAUNCH_MODE"] = "PARALLEL" 19 | os.environ["NCCL_IB_HCA"] = subprocess.getoutput( 20 | "pushd /sys/class/infiniband/ > /dev/null; for i in mlx5_*; " 21 | "do cat $i/ports/1/gid_attrs/types/* 2>/dev/null " 22 | "| grep v >/dev/null && echo $i ; done; popd > /dev/null" 23 | ) 24 | os.environ["NCCL_IB_GID_INDEX"] = "3" 25 | os.environ["NCCL_IB_TC"] = "106" 26 | 27 | 28 | def configure_omp(num_threads=1): 29 | """ 30 | If OMP_NUM_THREADS is not configured and world_size is greater than 1, 31 | Configure OMP_NUM_THREADS environment variables of NCCL to `num_thread`. 32 | 33 | Args: 34 | num_threads (int): value of `OMP_NUM_THREADS` to set. 35 | """ 36 | # We set OMP_NUM_THREADS=1 by default, which achieves the best speed on our machines 37 | # feel free to change it for better performance. 38 | if "OMP_NUM_THREADS" not in os.environ and get_world_size() > 1: 39 | os.environ["OMP_NUM_THREADS"] = str(num_threads) 40 | if is_main_process(): 41 | logger.info( 42 | "\n***************************************************************\n" 43 | "We set `OMP_NUM_THREADS` for each process to {} to speed up.\n" 44 | "please further tune the variable for optimal performance.\n" 45 | "***************************************************************".format( 46 | os.environ["OMP_NUM_THREADS"] 47 | ) 48 | ) 49 | 50 | 51 | def configure_module(ulimit_value=8192): 52 | """ 53 | Configure pytorch module environment. setting of ulimit and cv2 will be set. 54 | 55 | Args: 56 | ulimit_value(int): default open file number on linux. Default value: 8192. 57 | """ 58 | # system setting 59 | try: 60 | import resource 61 | 62 | rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) 63 | resource.setrlimit(resource.RLIMIT_NOFILE, (ulimit_value, rlimit[1])) 64 | except Exception: 65 | # Exception might be raised in Windows OS or rlimit reaches max limit number. 66 | # However, set rlimit value might not be necessary. 67 | pass 68 | 69 | # cv2 70 | # multiprocess might be harmful on performance of torch dataloader 71 | os.environ["OPENCV_OPENCL_RUNTIME"] = "disabled" 72 | try: 73 | cv2.setNumThreads(0) 74 | cv2.ocl.setUseOpenCL(False) 75 | except Exception: 76 | # cv2 version mismatch might rasie exceptions. 77 | pass 78 | --------------------------------------------------------------------------------