From 513a78b44272e75fd654b29b16000cd718705c4f Mon Sep 17 00:00:00 2001
From: zyayoung <1328410180@qq.com>
Date: Fri, 11 Feb 2022 21:53:40 +0800
Subject: [PATCH] results on dance track and minor fix

---
 .gitignore        |  8 ++++++++
 README.md         | 15 +++++++++++++--
 datasets/dance.py |  2 +-
 models/motr.py    |  1 +
 submit_dance.py   |  4 ----
 5 files changed, 23 insertions(+), 7 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..559c6ce
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,8 @@
+__pycache__/
+*.pth
+*.train
+exps/
+build/
+*.egg
+*.egg-info
+*.mp4
diff --git a/README.md b/README.md
index 00d6538..431b148 100644
--- a/README.md
+++ b/README.md
@@ -23,6 +23,8 @@ This repository is an official implementation of the paper [MOTR: End-to-End Mul
 ## Updates
 - (2021/09/23) Report BDD100K results and release corresponding codes [motr_bdd100k](https://github.com/megvii-model/MOTR/tree/motr_bdd100k). 
 - (2022/02/09) Higher performance achieved by not clipping the bounding boxes inside the image.
+- (2022/02/11) Add checkpoint support for training on RTX 2080ti.
+- (2022/02/11) Report [DanceTrack](https://github.com/DanceTrack/DanceTrack) results and [scripts](configs/r50_motr_train_dance.sh).
 
 ## Main Results
 
@@ -32,6 +34,12 @@ This repository is an official implementation of the paper [MOTR: End-to-End Mul
 | :--------: | :---------: | :------------------: | :------: | :------: | :-----: | :------: | :------: | :------: | :-----------------------------------------------------------------------------------------: |
 |    MOTR    |    MOT17    | MOT17+CrowdHuman Val |   71.1   |   68.4   |  2229   |   56.9   |   55.8   |   58.4   | [model](https://drive.google.com/file/d/1Utd6aqnuuOiMSQGvq4UGsnAfMpjtsW3E/view?usp=sharing) |
 
+### DanceTrack
+
+| **Method** | **Dataset** | **Train Data** | **MOTA** | **IDF1** | **HOTA** | **AssA** | **DetA** |                                           **URL**                                           |
+| :--------: | :---------: | :------------: | :------: | :------: | :------: | :------: | :------: | :-----------------------------------------------------------------------------------------: |
+|    MOTR    | DanceTrack  |   DanceTrack   |   79.7   |   51.5   |   54.2   |   40.2   |   73.5   | [model](https://drive.google.com/file/d/1zs5o1oK8diafVfewRl3heSVQ7-XAty3J/view?usp=sharing) |
+
 ### BDD100K
 
 | **Method** | **Dataset** | **Train Data** | **MOTA** | **IDF1** | **IDS** |                                           **URL**                                           |
@@ -40,8 +48,8 @@ This repository is an official implementation of the paper [MOTR: End-to-End Mul
 
 *Note:*
 
-1. All models of MOTR are trained on 8 NVIDIA Tesla V100 GPUs.
-2. The training time is about 2.5 days for 200 epochs;
+1. MOTR on MOT17 and DanceTrack is trained on 8 NVIDIA RTX 2080ti GPUs.
+2. The training time for MOT17 is about 2.5 days on V100 or 4 days on RTX 2080ti;
 3. The inference speed is about 7.5 FPS for resolution 1536x800;
 4. All models of MOTR are trained with ResNet50 with pre-trained weights on COCO dataset.
 
@@ -102,6 +110,9 @@ The codebase is built on top of [Deformable DETR](https://github.com/fundamental
 ├── MOT17
 │   ├── images
 │   ├── labels_with_ids
+├── DanceTrack
+│   ├── train
+│   ├── test
 ├── bdd100k
 │   ├── images
 │       ├── track
diff --git a/datasets/dance.py b/datasets/dance.py
index 834fd1b..cf80a18 100644
--- a/datasets/dance.py
+++ b/datasets/dance.py
@@ -253,5 +253,5 @@ def build(image_set, args):
         dataset = DetMOTDetection(args, data_txt_path=data_txt_path, seqs_folder=root, dataset2transform=dataset2transform)
     if image_set == 'val':
         data_txt_path = args.data_txt_path_val
-        dataset = DetMOTDetection(args, data_txt_path=data_txt_path, seqs_folder=root, dataset2transform=dataset2transform, joint=False)
+        dataset = DetMOTDetection(args, data_txt_path=data_txt_path, seqs_folder=root, dataset2transform=dataset2transform)
     return dataset
diff --git a/models/motr.py b/models/motr.py
index a964987..c358e3f 100644
--- a/models/motr.py
+++ b/models/motr.py
@@ -659,6 +659,7 @@ def build(args):
         'coco': 91,
         'coco_panoptic': 250,
         'e2e_mot': 1,
+        'e2e_dance': 1,
         'e2e_joint': 1,
         'e2e_static_mot': 1,
     }
diff --git a/submit_dance.py b/submit_dance.py
index b4e39d8..1488f70 100644
--- a/submit_dance.py
+++ b/submit_dance.py
@@ -499,10 +499,6 @@ def detect(self, prob_threshold=0.7, area_threshold=100, vis=False):
     sub_dir = 'DanceTrack/test'
     seq_nums = os.listdir(os.path.join(args.mot_path, sub_dir))
 
-    rank = int(os.environ.get('RLAUNCH_REPLICA', '0'))
-    ws = int(os.environ.get('RLAUNCH_REPLICA_TOTAL', '1'))
-    seq_nums = seq_nums[rank::ws]
-
     for seq_num in seq_nums:
         det = Detector(args, model=detr, seq_num=seq_num)
         det.detect()