Implement TDSP - DanceTrack SOTA

configs/appearance.yaml

@@ -2,10 +2,10 @@ defaults:
   - the_global_config
   - resources: default.yaml
   - dataset: dancetrack_appearance.yaml
-  - train: id.yaml
+  - train: bce.yaml
   - eval: default.yaml
-  - model_config: mm_appearance.yaml
+  - model_config: mm_tdsp_appearance.yaml
   - path: default.yaml
 
-experiment_name: exp74a-fromExp73-HalvedLrDoubleEpochs
+experiment_name: exp107a-fromExp105-ScaleUp
 dataset_name: DanceTrack

configs/bbox_only.yaml → configs/bbox.yaml

@@ -2,10 +2,10 @@ defaults:
   - the_global_config
   - resources: default.yaml
   - dataset: dancetrack_bbox.yaml
-  - train: batch.yaml
+  - train: bce.yaml
   - eval: default.yaml
-  - model_config: mm_bboxes.yaml
+  - model_config: mm_tdsp_bboxes.yaml
   - path: default.yaml
 
-experiment_name: exp74b-fromExp73-HalvedLrDoubleEpochs
+experiment_name: exp108b-fromExp107-BboxEmbDim512
 dataset_name: DanceTrack

configs/dataset/augmentations/default.yaml

@@ -19,7 +19,7 @@ augmentations:
   - _target_: mot_jepa.datasets.dataset.augmentations.video.IdentitySwitchAugmentation
     switch_ratio: 0.3
   - _target_: mot_jepa.datasets.dataset.augmentations.appearance.AppearanceNoiseAugmentation
-    alpha: 0.5
+    alpha: 0.40
   - _target_: mot_jepa.datasets.dataset.augmentations.video.SmartIdentitySwitchAugmentation
     switch_ratio: 0.5
     iou_threshold: 0.5

configs/dataset/dancetrack.yaml

@@ -2,7 +2,7 @@ defaults:
   - transform: scaled_bbox_keypoints.yaml
   - augmentations: default.yaml
   - feature_extractor: pred_bbox_keypoints_appearance.yaml
-  - sampler: scene_sampler.yaml
+  # - sampler: scene_sampler_within_range.yaml
 
 index:
   type: mot
@@ -11,10 +11,10 @@ index:
       - /media/home/DanceTrack-orig/
 
 n_tracks: 40
-clip_length: 30
+clip_length: 50
 min_clip_tracks: 1
 clip_sampling_step: 1
 val_clip_sampling_step: 1
 
 sampler: null
-use_batch_sampler: false
+use_batch_sampler: false

configs/dataset/dancetrack_appearance.yaml

@@ -2,6 +2,7 @@ defaults:
   - transform: scaled_bbox_keypoints.yaml
   - augmentations: default.yaml
   - feature_extractor: pred_appearance.yaml
+  # - sampler: scene_sampler_within_range.yaml
 
 index:
   type: mot

configs/dataset/feature_extractor/pred_appearance.yaml

@@ -3,4 +3,7 @@ extractor_params:
   prediction_path: /media/home/cameltrack-states/extracted-features
   extra_false_positives: true
   feature_names:
-    - appearance
+    - appearance
+
+  random_appearance_jitter_ratio: 0.0
+  random_appearance_jitter_range: 0

configs/dataset/sampler/scene_sampler.yaml

@@ -1,3 +1,3 @@
 _target_: 'mot_jepa.datasets.dataset.sampler.scene_sampler.SceneBatchSamplerWithRepeat.from_dataset'
-n_scenes: 4
-n_frames: 8
+n_scenes: 12
+n_frames: 1

configs/dataset/sampler/scene_sampler_within_range.yaml

@@ -0,0 +1,3 @@
+_target_: 'mot_jepa.datasets.dataset.sampler.scene_sampler.OneSceneWithRangeSampler.from_dataset'
+n_scenes: 5
+n_frames: 2

configs/default.yaml

@@ -2,10 +2,25 @@ defaults:
   - the_global_config
   - resources: default.yaml
   - dataset: dancetrack.yaml
-  - train: mm.yaml
+  - train: bce.yaml
   - eval: default.yaml
-  - model_config: mm_bboxes_keypoints_appearance.yaml
+  - model_config: mm_tdsp_bboxes_keypoints_appearance.yaml
   - path: default.yaml
 
-experiment_name: exp74-fromExp73-HalvedLrDoubleEpochs
+experiment_name: exp111-fromExp110-LinearSumDecoder
 dataset_name: DanceTrack
+
+resources:
+  batch_size: 8
+
+train:
+  max_epochs: 10
+
+  optimizer_config:
+    _target_: torch.optim.AdamW
+    lr: 1e-5
+    weight_decay: 1e-3
+
+  scheduler_config:
+    _target_: mot_jepa.trainer.scheduler.create_warmup_cosine_annealing_scheduler
+    n_warmup_epochs: 1

configs/keypoints.yaml

@@ -2,10 +2,10 @@ defaults:
   - the_global_config
   - resources: default.yaml
   - dataset: dancetrack_keypoints.yaml
-  - train: batch.yaml
+  - train: bce.yaml
   - eval: default.yaml
-  - model_config: mm_keypoints.yaml
+  - model_config: mm_tdsp_keypoints.yaml
   - path: default.yaml
 
-experiment_name: exp74k-fromExp73-HalvedLrDoubleEpochs
+experiment_name: exp108k-fromExp107-BboxEmbDim512
 dataset_name: DanceTrack

configs/model_config/mm_appearance.yaml

@@ -7,7 +7,7 @@ common_params:
   track_encoder_n_layers: 2
   track_encoder_ffn_dim: 512
   projector_intermediate_dim: 512
-  interaction_encoder_enable: true
+  interaction_encoder_enable: false
   interaction_encoder_n_heads: 8
   interaction_encoder_n_layers: 2
   interaction_encoder_ffn_dim: 512

configs/model_config/mm_bboxes.yaml

@@ -7,7 +7,7 @@ common_params:
   track_encoder_n_layers: 2
   track_encoder_ffn_dim: 512
   projector_intermediate_dim: 512
-  interaction_encoder_enable: true
+  interaction_encoder_enable: false
   interaction_encoder_n_heads: 8
   interaction_encoder_n_layers: 2
   interaction_encoder_ffn_dim: 512

configs/model_config/mm_bboxes_keypoints.yaml

@@ -6,7 +6,7 @@ common_params:
   track_encoder_n_layers: 2
   track_encoder_ffn_dim: 512
   projector_intermediate_dim: 512
-  interaction_encoder_enable: true
+  interaction_encoder_enable: false
   interaction_encoder_n_heads: 8
   interaction_encoder_n_layers: 2
   interaction_encoder_ffn_dim: 512

configs/model_config/mm_bboxes_keypoints_appearance.yaml

@@ -7,7 +7,7 @@ common_params:
   track_encoder_n_layers: 2
   track_encoder_ffn_dim: 512
   projector_intermediate_dim: 512
-  interaction_encoder_enable: true
+  interaction_encoder_enable: false
   interaction_encoder_n_heads: 8
   interaction_encoder_n_layers: 2
   interaction_encoder_ffn_dim: 512
@@ -35,9 +35,9 @@ aggregator_params:
   dropout: 0.1
 
 per_feature_checkpoint:
-  bbox: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp74b-fromExp73-HalvedLrDoubleEpochs/checkpoints/last.pt
-  keypoints: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp74k-fromExp73-HalvedLrDoubleEpochs/checkpoints/last.pt
-  appearance: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp74a-fromExp73-HalvedLrDoubleEpochs/checkpoints/last.pt
+  bbox: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp93b-fromExp90-ClipLevelBCE/checkpoints/last.pt
+  keypoints: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp93k-fromExp90-ClipLevelBCE/checkpoints/last.pt
+  appearance: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp98a-fromExp97-EmbDim512/checkpoints/last.pt
 
 object_interaction_encoder_enable: true
 object_interaction_encoder_params:

configs/model_config/mm_keypoints.yaml

@@ -7,7 +7,7 @@ common_params:
   track_encoder_n_layers: 2
   track_encoder_ffn_dim: 512
   projector_intermediate_dim: 512
-  interaction_encoder_enable: true
+  interaction_encoder_enable: false
   interaction_encoder_n_heads: 8
   interaction_encoder_n_layers: 2
   interaction_encoder_ffn_dim: 512

configs/model_config/mm_tdsp_appearance.yaml

@@ -0,0 +1,33 @@
+_target_: mot_jepa.architectures.tdcp.core.build_mm_tdsp_model
+mm_dim: 1024
+similarity_prediction_head_hidden_dim: 512
+
+sph_common_params:
+  hidden_dim: 512
+sph_per_feature_params:
+  appearance:
+    hidden_dim: 512
+
+common_params:
+  hidden_dim: 256
+  dropout: 0.1
+  track_encoder_n_heads: 8
+  track_encoder_n_layers: 4
+  track_encoder_ffn_dim: 512
+  projector_intermediate_dim: 512
+  interaction_encoder_enable: true
+  interaction_encoder_n_heads: 8
+  interaction_encoder_n_layers: 4
+  interaction_encoder_ffn_dim: 512
+per_feature_params:
+  appearance:
+    hidden_dim: 512
+    feature_encoder_type: parts_appearance
+    feature_encoder_params:
+      emb_size: 128
+      hidden_dim: 512
+    track_encoder_enable_motion_encoder: false
+    track_encoder_ffn_dim: 1024
+    interaction_encoder_ffn_dim: 1024
+aggregator_type: sum
+aggregator_params: {}

configs/model_config/mm_tdsp_bboxes.yaml

@@ -0,0 +1,28 @@
+_target_: mot_jepa.architectures.tdcp.core.build_mm_tdsp_model
+mm_dim: 1024
+similarity_prediction_head_hidden_dim: 512
+
+sph_common_params:
+  hidden_dim: 512
+sph_per_feature_params:
+  bbox:
+    hidden_dim: 512
+
+common_params:
+  hidden_dim: 512
+  dropout: 0.1
+  track_encoder_n_heads: 8
+  track_encoder_n_layers: 4
+  track_encoder_ffn_dim: 1024
+  projector_intermediate_dim: 512
+  interaction_encoder_enable: true
+  interaction_encoder_n_heads: 8
+  interaction_encoder_n_layers: 4
+  interaction_encoder_ffn_dim: 1024
+per_feature_params:
+  bbox:
+    feature_encoder_type: motion
+    feature_encoder_params:
+      input_dim: 5
+aggregator_type: sum
+aggregator_params: {}

configs/model_config/mm_tdsp_bboxes_keypoints_appearance.yaml

@@ -0,0 +1,61 @@
+_target_: mot_jepa.architectures.tdcp.core.build_mm_tdsp_model
+mm_dim: 1024
+similarity_prediction_head_hidden_dim: 1024
+
+sph_common_params:
+  hidden_dim: 512
+sph_per_feature_params:
+  bbox:
+    hidden_dim: 512
+  keypoints:
+    hidden_dim: 512
+  appearance:
+    hidden_dim: 512
+
+common_params:
+  hidden_dim: 512
+  dropout: 0.1
+  track_encoder_n_heads: 8
+  track_encoder_n_layers: 4
+  track_encoder_ffn_dim: 1024
+  projector_intermediate_dim: 512
+  interaction_encoder_enable: true
+  interaction_encoder_n_heads: 8
+  interaction_encoder_n_layers: 4
+  interaction_encoder_ffn_dim: 1024
+per_feature_params:
+  bbox:
+    feature_encoder_type: motion
+    feature_encoder_params:
+      input_dim: 5
+  keypoints:
+    feature_encoder_type: motion
+    feature_encoder_params:
+      input_dim: 35
+  appearance:
+    hidden_dim: 512
+    feature_encoder_type: parts_appearance
+    feature_encoder_params:
+      emb_size: 128
+      hidden_dim: 512
+    track_encoder_enable_motion_encoder: false
+    track_encoder_ffn_dim: 1024
+    interaction_encoder_ffn_dim: 1024
+aggregator_type: sum # transformer
+aggregator_params: {}
+  # hidden_dim: ${model_config.mm_dim}
+  # n_heads: 8
+  # n_layers: 4
+  # dropout: 0.1
+
+per_feature_checkpoint:
+  bbox: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp108b-fromExp107-BboxEmbDim512/checkpoints/last.pt
+  keypoints: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp108k-fromExp107-BboxEmbDim512/checkpoints/last.pt
+  appearance: /media/home/MOT-JEPA-outputs/experiments/DanceTrack/exp107a-fromExp105-ScaleUp/checkpoints/last.pt
+
+object_interaction_encoder_enable: true
+object_interaction_encoder_params:
+  hidden_dim: ${model_config.mm_dim}
+  n_heads: 8
+  n_layers: 2
+  dropout: 0.1

configs/model_config/mm_tdsp_keypoints.yaml

@@ -0,0 +1,28 @@
+_target_: mot_jepa.architectures.tdcp.core.build_mm_tdsp_model
+mm_dim: 1024
+similarity_prediction_head_hidden_dim: 512
+
+sph_common_params:
+  hidden_dim: 512
+sph_per_feature_params:
+  keypoints:
+    hidden_dim: 512
+
+common_params:
+  hidden_dim: 512
+  dropout: 0.1
+  track_encoder_n_heads: 8
+  track_encoder_n_layers: 4
+  track_encoder_ffn_dim: 1024
+  projector_intermediate_dim: 512
+  interaction_encoder_enable: true
+  interaction_encoder_n_heads: 8
+  interaction_encoder_n_layers: 4
+  interaction_encoder_ffn_dim: 1024
+per_feature_params:
+  keypoints:
+    feature_encoder_type: motion
+    feature_encoder_params:
+      input_dim: 35
+aggregator_type: sum
+aggregator_params: {}

configs/resources/default.yaml

@@ -1,4 +1,4 @@
-batch_size: 8
-val_batch_size: 4
+batch_size: 32
+val_batch_size: 16
 accelerator: 'cuda:0'
-num_workers: 12
+num_workers: 20

configs/train/bce.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - base.yaml
+  - _self_
+
+loss_config:
+  _target_: mot_jepa.trainer.losses.bce.ClipLevelBCE
+  pos_weight: 10.0
+  assoc_threshold: 1e-2

configs/train/clip.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - base.yaml
+  - _self_
+
+loss_config:
+  _target_: mot_jepa.trainer.losses.infonce.ClipLevelInfoNCE

configs/train/mm.yaml

@@ -5,14 +5,14 @@ defaults:
 loss_config:
   _target_: mot_jepa.trainer.losses.infonce.MultiFeatureLoss
   mm_loss:
-    _target_: mot_jepa.trainer.losses.infonce.IDLevelInfoNCE
+    _target_: mot_jepa.trainer.losses.infonce.ClipLevelInfoNCE
   per_feature_losses:
     bbox:
-      _target_: mot_jepa.trainer.losses.infonce.BatchLevelInfoNCE
+      _target_: mot_jepa.trainer.losses.infonce.ClipLevelInfoNCE
     keypoints:
-      _target_: mot_jepa.trainer.losses.infonce.BatchLevelInfoNCE
+      _target_: mot_jepa.trainer.losses.infonce.ClipLevelInfoNCE
     appearance:
-      _target_: mot_jepa.trainer.losses.infonce.IDLevelInfoNCE
+      _target_: mot_jepa.trainer.losses.infonce.ClipLevelInfoNCE
   per_feature_weights:
     bbox: 0.3
     keypoints: 0.3