Add causal person fusion and MuVAP stage-3 baselines

third_party/EgoSpeak/MiniROAD/.gitignore

@@ -0,0 +1,27 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+
+# PyTorch models
+*.pth
+
+# Miscellaneous
+*.DS_Store
+*.swp
+*.pyc
+*.pkl
+
+*.out
+
+dataset/
+# Dataset soft links
+data/
+
+# Output files
+output/
+output_*/
+
+checkpoints/
+
+notebooks/
+wandb/

third_party/EgoSpeak/MiniROAD/configs/EasyCom/mamba_aa_easycom.yaml

@@ -0,0 +1,32 @@
+model: 'MambaA'
+data_name: 'EasyCom_ANTICIPATION'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/EasyCom'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2'
+annotation_type: 'target_perframe'
+video_list_path: 'data_info/data_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 64 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 3
+embedding_dim: 2048
+hidden_dim: 1024
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100
+seed: 20

third_party/EgoSpeak/MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml

@@ -0,0 +1,32 @@
+model: 'MiniROADA'
+data_name: 'EasyCom_ANTICIPATION'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/EasyCom'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2'
+annotation_type: 'target_perframe'
+video_list_path: 'data_info/data_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 64 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 3 # [background, target speaker speaking, other person speaking]
+embedding_dim: 2048
+hidden_dim: 1024
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100
+seed: 20

third_party/EgoSpeak/MiniROAD/configs/YTConv/mamba_aa_ytconv.yaml

@@ -0,0 +1,32 @@
+model: 'MambaA'
+data_name: 'YTConv_ANTICIPATION'
+data_info_name: 'shard_0'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/YTConv/shard_0'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2' # or 'none'?
+annotation_type: 'target_perframe/case1'
+video_list_path: 'data_info/YTConv_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 32 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 2 # for YTConv pretraining, there is no 
+embedding_dim: 2048
+hidden_dim: 1024
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100

third_party/EgoSpeak/MiniROAD/configs/YTConv/miniroad_aa_ytconv.yaml

@@ -0,0 +1,32 @@
+model: 'MiniROADA'
+data_name: 'YTConv_ANTICIPATION'
+data_info_name: 'shard_0'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/YTConv/shard_0'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2'
+annotation_type: 'target_perframe'
+video_list_path: 'data_info/YTConv_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 64 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 2 # for pretraining, there is no target speaker
+embedding_dim: 2048
+hidden_dim: 1024
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100

third_party/EgoSpeak/MiniROAD/configs/ego4d/mamba_aa_ego4dshuffle.yaml

@@ -0,0 +1,36 @@
+model: 'MambaA'
+data_name: 'EGO4D_ANTICIPATION'
+data_info_name: 'EGO4DSHUFFLE'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/EGO4D'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2'
+annotation_type: 'target_perframe'
+video_list_path: 'data_info/data_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 64 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 3 # including background
+embedding_dim: 2048
+hidden_dim: 1024
+d_state: 16
+d_conv: 4
+expand: 2 
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100
+seed: 20

third_party/EgoSpeak/MiniROAD/configs/ego4d/miniroad_aa_ego4dshuffle.yaml

@@ -0,0 +1,32 @@
+model: 'MiniROADA'
+data_name: 'EGO4D_ANTICIPATION'
+data_info_name: 'EGO4DSHUFFLE'
+task: 'ANTICIPATION'
+loss: 'ANTICIPATION' # NONUNIFORM, originally
+metric: 'AP' # AP or cAP
+optimizer: 'AdamW'
+device: 'cuda'
+feature_pretrained: 'resnet50'
+root_path: './data/EGO4D'
+rgb_type: 'rgb_kinetics_resnet50'
+flow_type: 'audio_wav2vec2' # or 'none'?
+annotation_type: 'target_perframe'
+video_list_path: 'data_info/data_info.json'
+output_path: './output_cfg'
+window_size: 128
+batch_size: 64 
+test_batch_size: 1 
+# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L
+num_epoch: 30
+lr: 0.0001
+weight_decay: 0.05
+num_workers: 4
+dropout: 0.20
+num_classes: 3
+embedding_dim: 2048
+hidden_dim: 1024
+num_layers: 1
+stride: 4
+anticipation_length: 10
+actionness: False
+ignore_index: -100

third_party/EgoSpeak/MiniROAD/configs/feature_sizes.py

@@ -0,0 +1,26 @@
+# dimension of features extracted from different models
+# used in feature_head, dataset, models
+
+FEATURE_SIZES = {
+    'rgb_anet_resnet50': 2048,
+    'flow_anet_resnet50': 2048,
+    'rgb_kinetics_bninception': 1024,
+    'flow_kinetics_bninception': 1024,
+    'rgb_farl': 512,
+    'rgb_kinetics_resnet50': 2048,
+    'rgb_kinetics_resnet50_debug': 2048,
+    'rgb_kinetics_resnet50_debug_resized': 2048,
+    'flow_kinetics_resnet50': 2048,
+    'flow_nv_kinetics_bninception': 1024,
+    'rgb_kinetics_i3d': 2048,
+    'flow_kinetics_i3d': 2048,
+    'audio_wav2vec2': 5120,
+    'audio_wav2vec2_left': 5120,
+    'audio_wav2vec2_right': 5120,
+    'text_word2vec': 300,
+    'text_fasttext_crawl': 300,
+    'text_gpt2_last_subword_window_size_0': 768,
+    'text_gpt2_last_subword_window_size_20': 768,
+    'text_gpt2_last_subword_window_size_50': 768,
+    'text_gpt2_last_subword_window_size_100': 768,
+}

third_party/EgoSpeak/MiniROAD/criterions/__init__.py

@@ -0,0 +1,2 @@
+from .loss import OadLoss
+from .loss_builder import build_criterion

third_party/EgoSpeak/MiniROAD/criterions/loss.py

@@ -0,0 +1,205 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from criterions.loss_builder import CRITERIONS
+
+
+@CRITERIONS.register('NONUNIFORM')
+class OadLoss(nn.Module):
+    
+    def __init__(self, cfg, reduction='mean'):
+        super(OadLoss, self).__init__()
+        self.reduction = reduction
+        self.num_classes = cfg['num_classes']
+        self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
+        assert self.loss_offset==0, f'loss_offset is not 0: {self.loss_offset}'
+        self.loss = self.end_loss
+
+    def end_loss(self, out_dict, target):
+        # logits: (B, seq, K) target: (B, seq, K)
+        assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}'
+        logits = out_dict['logits']
+        logits = logits[:,-1 - self.loss_offset,:].contiguous()
+        target = target[:,-1 - self.loss_offset,:].contiguous()
+        ce_loss = self.mlce_loss(logits, target)
+        return ce_loss
+
+    def mlce_loss(self, logits, target):
+        '''
+        multi label cross entropy loss. 
+        logits: (B, K) target: (B, K) 
+        '''
+        logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
+        output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+        return loss
+
+    def forward(self, out_dict, target): 
+        return self.loss(out_dict, target)
+
+
+@CRITERIONS.register('UNIFORM')
+class OadUniformLoss(nn.Module):
+    
+    def __init__(self, cfg, reduction='mean'):
+        super(OadUniformLoss, self).__init__()
+        self.reduction = reduction
+        self.num_classes = cfg['num_classes']
+        # self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
+        self.loss = self.end_loss
+
+    def end_loss(self, out_dict, target):
+        # logits: (B, seq, K) target: (B, seq, K)
+        logits = out_dict['logits']
+        # logits = logits[:,-1 - self.loss_offset,:].contiguous()
+        # target = target[:,-1 - self.loss_offset,:].contiguous()
+        ce_loss = self.mlce_loss(logits, target)
+        return ce_loss
+
+    def mlce_loss(self, logits, target):
+        '''
+        multi label cross entropy loss. 
+        logits: (B, K) target: (B, K) 
+        '''
+        logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
+        output = torch.sum(-target * logsoftmax(logits), dim=1) # B
+        # output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+        return loss
+
+    def forward(self, out_dict, target): 
+        return self.loss(out_dict, target)
+
+@CRITERIONS.register('OADTR')
+class OadTRLoss(nn.Module):
+    
+    def __init__(self, cfg, reduction='mean'):
+        super(OadTRLoss, self).__init__()
+        self.reduction = reduction
+        self.num_classes = cfg['num_classes']
+        # self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
+        self.loss = self.end_loss
+
+    def end_loss(self, out_dict, target):
+        # logits: (B, seq, K) target: (B, seq, K)
+        assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}'
+        logits = out_dict['logits']
+        # logits = logits[:,-1 - self.loss_offset,:].contiguous()
+        # target = target[:,-1 - self.loss_offset,:].contiguous()
+        ce_loss = self.mlce_loss(logits, target)
+        return ce_loss
+
+    def mlce_loss(self, logits, target):
+        '''
+        multi label cross entropy loss. 
+        logits: (B, K) target: (B, K) 
+        '''
+        logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
+        output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+        return loss
+
+    def forward(self, out_dict, enc_target, dec_target): 
+        return self.loss(out_dict, enc_target, dec_target)
+
+@CRITERIONS.register('ANTICIPATION')
+class OadAntLoss(nn.Module):
+    
+    # reduction was sum, not mean, what if we change to sum?
+    def __init__(self, cfg, reduction='mean'):
+        super(OadAntLoss, self).__init__()
+        self.reduction = reduction
+        self.loss = self.anticipation_loss
+        self.num_classes = cfg['num_classes']
+        
+    def anticipation_loss(self, out_dict, target, ant_target):
+        # anticipation_logits: (B, seq, ANTICIPATION_LENGTH, num_class)
+        # ant_target: (B, ANTICIPATION_LENGTH, num_class)
+        anticipation_logits = out_dict['anticipation_logits']
+        # [B, window_size, anticipation_length, num_class]
+        pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes)
+        
+        anticipation_logit_targets = ant_target.view(-1, self.num_classes)
+        ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets)
+
+        return ant_loss
+
+    def ce_loss(self, out_dict, target):
+        # logits: (B, seq, K) target: (B, seq, K)
+        logits = out_dict['logits']
+        logits = logits[:,-1,:].contiguous()
+        target = target[:,-1,:].contiguous()
+        ce_loss = self.mlce_loss(logits, target)
+        return ce_loss
+
+    def mlce_loss(self, logits, target):
+        '''
+        multi label cross entropy loss. 
+        logits: (B, K) target: (B, K) 
+        '''
+        logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
+        output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+
+        return loss
+
+    def forward(self, out_dict, target, ant_target): 
+        return self.loss(out_dict, target, ant_target)
+
+
+
+@CRITERIONS.register('UNIFORM_ANTICIPATION')
+class OadUniformAntLoss(nn.Module):
+    
+    # reduction was sum, not mean, what if we change to sum?
+    def __init__(self, cfg, reduction='mean'):
+        super(OadUniformAntLoss, self).__init__()
+        self.reduction = reduction
+        self.loss = self.anticipation_loss
+        self.num_classes = cfg['num_classes']
+
+    def anticipation_loss(self, out_dict, target, ant_target):
+        anticipation_logits = out_dict['anticipation_logits']
+        pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes)
+        anticipation_logit_targets = ant_target.view(-1, self.num_classes)
+        ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets)
+        return ant_loss
+
+    def ce_loss(self, out_dict, target):
+        # logits: (B, seq, K) target: (B, seq, K)
+        logits = out_dict['logits']
+        logits = logits[:,-1,:].contiguous()
+        target = target[:,-1,:].contiguous()
+        ce_loss = self.mlce_loss(logits, target)
+        return ce_loss
+
+    def mlce_loss(self, logits, target):
+        '''
+        multi label cross entropy loss. 
+        logits: (B, K) target: (B, K) 
+        '''
+        logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
+        output = torch.sum(-target * logsoftmax(logits), dim=1) # B
+        
+        # output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+
+        return loss
+
+    def forward(self, out_dict, target, ant_target): 
+        return self.loss(out_dict, target, ant_target)

third_party/EgoSpeak/MiniROAD/criterions/loss_builder.py

@@ -0,0 +1,11 @@
+__all__ = [
+    'build_criterion'
+]
+
+from utils import Registry
+
+CRITERIONS = Registry()
+
+def build_criterion(cfg, device=None):
+    criterion = CRITERIONS[cfg["loss"]](cfg)
+    return criterion.to(device)

third_party/EgoSpeak/MiniROAD/data_info/YTConv_info.json

@@ -0,0 +1,426 @@
+{
+    "shard_0": {
+        "class_index": [
+            "background",
+            "camera wearer start to speak"
+        ],
+        "train_session_set": [
+            "m5Aoywf51L4",
+            "p99cIB7nHjw",
+            "zGgMzBfPfaI",
+            "O0eCtKU3FxY",
+            "goZQ6iOIFYo",
+            "9Tb-lSyhgoE",
+            "GBxQtGUjAMA",
+            "VXgt7keOC8I",
+            "hZ2tin2J0P4",
+            "Ur-P18sm4tI",
+            "WfYmsNBQwg8",
+            "AvMNaDDfrN0",
+            "fW7pCN3wwdo",
+            "xG6HBdTtS4s",
+            "HcDlqUeDD9Q",
+            "1DE3mKncHUA",
+            "oEMX_quSMZY",
+            "jmS3nWbGV5Y",
+            "qUJRjPRH-u4",
+            "aHsUgWfUcss",
+            "ZfdQOmJkaK0",
+            "4f43nYrfTnI",
+            "qE1FnzkNjrE",
+            "4NsmKhs9flk",
+            "3wDSJ9ABzmo",
+            "sY6_8QwRZsw",
+            "hByfAzEoeKA",
+            "vDEsV0AhcjM",
+            "XGNhX3Nu1RE",
+            "MkVT8K0V57A",
+            "yR8y-wN8Clk",
+            "P1dTOlO2sl8",
+            "6UGC_wxTeB0",
+            "d4Xud2Gqoxk",
+            "IJ1aqePculQ",
+            "o8Oxf5QBitQ",
+            "XVci_V4qxmc",
+            "ofDzhOS1Pv0",
+            "mygTXBuMRF0",
+            "iq6q82Jo8ZQ",
+            "Kzc_ekBJ4QE",
+            "WEeJ9B1i8HI",
+            "rciqpxdCc9Y",
+            "Rh7d2D7JVLI",
+            "Qkx4hSlc8jk",
+            "ixi7sVLFrWM",
+            "n96-scSyoOQ",
+            "YZ5xYioIQdA",
+            "X76QxuWm1oA",
+            "2I7d_LfK75Y",
+            "AeUWBaYFMyo",
+            "lsqXFgymLlI",
+            "Dx_0jYMWN5o",
+            "0raJ1-1BJ7s",
+            "EGX8am5hWh0",
+            "KjtZ_CLY1hQ",
+            "R4FQI2BI3hw",
+            "fAoqqtsNjbQ",
+            "l6RwyA6jUdM",
+            "GJW6dlagRnU",
+            "j_FGvmbp6aM",
+            "V2O15h1VrH8",
+            "P3c70VbDP3E",
+            "epAwpBdUiYU",
+            "6dpv_JeyL8w",
+            "dUKYxsvnuBM",
+            "L2_zn7h8sOQ",
+            "m09qdCkAhmY",
+            "7BB71tIae7A",
+            "6go4oViEkdc",
+            "MYimxshUK1c",
+            "v4eZXWcIRQU",
+            "vGI_w1Ml6S0",
+            "BeFndnkT00c",
+            "9JHRbL1jmq4",
+            "0rBaR5dd-Fs",
+            "EvhkIS_-K-w",
+            "pYVbLx-h5HI",
+            "fVciOXYDlqg",
+            "zWFH4N-CmDE",
+            "Aw4cauzfX5E",
+            "lxVpcQHtcWs",
+            "32dkF-aXgug",
+            "_USorgWOsUk",
+            "1yVMU93nmv8",
+            "aYkF4Ybk-l8",
+            "Vwt18p0BqzM",
+            "oOQM5Pvjut4",
+            "_WWSMNtNYx0",
+            "zzlhwg0IJ5Y",
+            "lJhIffkUhr8",
+            "PbS9nSH3mqM",
+            "8EGVgiFpwhg",
+            "5LhXpv5IFO4",
+            "ibD4APVzWgg",
+            "2nhcSQmopO4",
+            "gri79OfCRvo",
+            "lOwQN3QcRzE",
+            "wQyluRNtqFo",
+            "IZyx8zeiKB0",
+            "mZyNfKyzDHU",
+            "A9eGJWiuIKs",
+            "d37DZogot1w",
+            "7ZUSVabB6N4",
+            "m62GQXq-358",
+            "ArwPVsCZo2c",
+            "QK-BFkpO8s4",
+            "PgfX9bvmEhs",
+            "l1jTUhwWJYA",
+            "S1PGHO4evaI",
+            "M6uPJa5B8Ns",
+            "XXzz19J-1w8",
+            "U4uV1xmUaB0",
+            "skQ0OBpkKqY",
+            "mfW57TDpMtc",
+            "NAwKA75cDwE",
+            "eIEBI2_Y14Q",
+            "eWQGYpSoiYU",
+            "3DPIUWhYtn8",
+            "_elvaRJrhE0",
+            "PAA8IFhWJzM",
+            "1XDKqhq6bb0",
+            "ngi9gy3DzRY",
+            "Y0yONfuimeE",
+            "g9_cAgHa_4Q",
+            "BIG7jEw-j4w",
+            "TFdom0lkd-0",
+            "2zV2N8_DwWg",
+            "kTz7XB_jpB4",
+            "GuoglIuhNjY",
+            "L98U3mO38gQ",
+            "UfWDoRAzpC8",
+            "W6TakaXCIIQ",
+            "i6j1I6pha-M",
+            "issHFC8Gbio",
+            "4EzGoRkd6SI",
+            "QoS0GgF4sI4",
+            "A8hJg3D_TGY",
+            "A3N10suCRCw",
+            "6962tnWRsNY",
+            "lS4p8osXSJs",
+            "wbZ7ZhcZd6o",
+            "fsIvQvIlKwc",
+            "4jQSOSi2DA8",
+            "Vpt6tZBvXQU",
+            "kxOG3kpLmYw",
+            "Mzjlz0nq9AA",
+            "80cuXt2VSko",
+            "VvswhogSiY0",
+            "5Q6ky5VrOIA",
+            "I_XtcUpNRF8",
+            "QmtmTGycK9A",
+            "qp3x8iy7vkA",
+            "xGU8QDtPrQI",
+            "uWWUfqth_BA",
+            "WbM7w8r4VGQ",
+            "bt6Ed1qjulw",
+            "7vymT27scoY",
+            "b9AdV6oVDVg",
+            "nf2zqYUAOwg",
+            "A-LNJc-1Jc0",
+            "yCI-hpPiDVc",
+            "Uu5DME4FqIM",
+            "symjuH9iOTc",
+            "mPf2MTnTBc0",
+            "NDTlQ1KGvH8",
+            "g9vuDzFQSRU",
+            "B6iqnvleve0",
+            "9bbszAD7JL0",
+            "RIyMAuz7Ckg",
+            "mlrQHtBRMhc",
+            "3gaidlW7vig",
+            "DGdV97q9ndk",
+            "qthjlg0lXFg",
+            "gWng5YLPOow",
+            "iaYG-KL9XRM",
+            "385414AVZcA",
+            "7xhBnWGpI7c",
+            "Nh0cmSQUs-E",
+            "Rvtj0QHPulA",
+            "rAtLksYQukw",
+            "S9rB__zbbe0",
+            "SNpMVsOPfHM",
+            "zhF43F4kO2w",
+            "HOUkcInyxN0",
+            "x8ymkfr2Z1o",
+            "Ix6Khtq6O-k",
+            "VHs_JyPhmtE",
+            "zJX3byjYSMc",
+            "zdqX09zZNf8",
+            "uPwlmwNQq8Y",
+            "e4xNTXrk6wE",
+            "P2PNUWN4GMs",
+            "_WVHCYjHBVY",
+            "5Jtv1GpPvHM",
+            "AIZYucKOzhw",
+            "vbsR_VDZXh0",
+            "9Ew0lDHB0kM",
+            "ejpkPVZHuMM",
+            "47HlvCfubE4",
+            "HzPkQSIG2m0",
+            "e7mTRnNnOXg",
+            "eZHx_biEECg",
+            "FvP6WFKNP4Y",
+            "g_TJ5XFEVfY",
+            "8Utg2VEWsJg",
+            "VMdBjL-tpoo",
+            "zihP2Eqeqxk",
+            "ILVqgOlntLM",
+            "Z1BWQI8KPCg",
+            "j8TRVp-7l3Q",
+            "G5gJBZ-3OlY",
+            "o1JUlGSirSE",
+            "OsPdhUQ0snE",
+            "-5uIaW2gKWQ",
+            "03kiFNh-9ME",
+            "w5W6mgDaJJA",
+            "sYsgd_8AFoo",
+            "vLg8kVUTmWU",
+            "OK0J0xsuSWQ",
+            "d284rjoU4Cs",
+            "5lE6V4jrUl4",
+            "bS1OpC9xu48",
+            "ZUxpqU_J_n0",
+            "AUliC0c_YIQ",
+            "nJ58-d_viOo",
+            "uIy3b2kzlrM",
+            "vOTLfb0qjfY",
+            "37_7tPZjC34",
+            "MUwCdP5qceI",
+            "lYvNi_ewd8U",
+            "PDG7SzDBTC4",
+            "QHQzoHolWvI",
+            "YmvOEj00LB8",
+            "fLhkXVx7Yw0",
+            "Uj5nhugVgwo",
+            "aHbrDu2o0WY",
+            "yKW6vVgEPCM",
+            "Q0WUgyqQ-4I",
+            "xw3KOPJ9a3E",
+            "5mrTo7c3bCM",
+            "x0ft032W--4",
+            "UdQxEbrIocc",
+            "EHZHAFyv_Ks",
+            "7DojKgPRLVk",
+            "-1CahRD6lJ0",
+            "f3TuYwl7AZE",
+            "TgGTovCcS9Y",
+            "AwyrJGeIN_g",
+            "RbPJD4Wh-l8",
+            "58bFdwYkDzU",
+            "6gIEKGxL7zo",
+            "BeuzJwTuChQ",
+            "YebZyAzLZuc",
+            "9T74b8EQBVY",
+            "TqSVFjw1uNM",
+            "WUYhCmfE1a0",
+            "vRCVUQFUio4",
+            "w2t9JdMNMgQ",
+            "nn6-x9xXMeE",
+            "FDFVPzLlzFs",
+            "qFhvi4bH29E",
+            "qdpsMXLsUZM",
+            "JREk1QeSc9A",
+            "jOVEs-4r1jc",
+            "v4VZI9Ouzy0",
+            "DMg5UUvf-C4",
+            "26J5zRzQkOg",
+            "YkiJ4iJfJbw",
+            "_d8j8nfTEIk",
+            "wMNb3hQXEUE",
+            "1KI8ex22blk",
+            "eCp2Jp7MtKg",
+            "HNlpHBb5dC4",
+            "C6_pDRdIo1k",
+            "ZRWJJHrgP54",
+            "L7jfx4veKIA",
+            "p1a5_qltICM",
+            "nZ0svM8DoY0",
+            "miEzjCJVaZk",
+            "TpC2tDCxABw",
+            "SZmqoqCX__U",
+            "dZ0bYteaI2Q",
+            "GW-cYEzShmY",
+            "iUYuPhutKVw",
+            "97VSm6T2zMw",
+            "jPiQyi5jr1M",
+            "lrUi4nsPZ9M",
+            "rXLeXke-k_w",
+            "mEO4izUqZ7E",
+            "4qI1LjMXfh0",
+            "mVt6myGMa-0",
+            "VJrnhh6FqEg",
+            "vwCeBWbZvrE",
+            "c08fl8-xmVE",
+            "syZ3bvXIrpA",
+            "i3ysNaulop4",
+            "__t3rU0QjsE",
+            "jT32uQQVGvU",
+            "Fzvmh-R2EWk",
+            "QApJSKrLg6E",
+            "L5y3j8SdGSU",
+            "3t8JGblR__0",
+            "e1F90W40t6g",
+            "j6L8UIa5z28",
+            "Zmg2_C1FuVI",
+            "XIXKJXsiBFc",
+            "NVHB1VdWFo4",
+            "-x7-8f7bsl4",
+            "2FSxSiowkGI",
+            "jgFIBIeFve0",
+            "1BLuXK1M5CE",
+            "zmpx6LUT2ow",
+            "z-QUzJTsFro",
+            "0sS4v13DAK8",
+            "gMIBQ0mbRvM",
+            "9kcaEXlTB5s",
+            "ZYArPGMLS4k",
+            "tejW9GMaucA",
+            "z9r4bmXCWiI",
+            "4Zpsb46iaLc",
+            "52Odc8UbvC8",
+            "x-B6-iYDHPs",
+            "8GinSDc0Is8",
+            "5lwbp5d_w0Y",
+            "I4LpTQLzlHE",
+            "lMl6-XH7ozY",
+            "tOubscNbY1Q",
+            "7gb7AH-oWcQ",
+            "mDerqt9Ib-I",
+            "gKV3FNd9WHE",
+            "kH5L1E_r-3U",
+            "sq5TRFxVp0w",
+            "SrBhj7Xusk8",
+            "RxlxEapsz2E",
+            "aCD3_TziiKo",
+            "_HPVZOkxdkk",
+            "pNCjgVQ53Y8",
+            "iPdT5GUoGvA",
+            "cIrlet1oRw4",
+            "1dyWUMBjMwQ",
+            "XZQB3cq18yc",
+            "i75Mn_AeFBY",
+            "_0VwR9WPS-k",
+            "q1JDjgnRdGY",
+            "1JNEpfvRq0Y",
+            "2FGj6HciQZA",
+            "eLi59TtNwHY",
+            "Z3W8OST7tmA",
+            "F2dNYov8-xw",
+            "kPfyJu_CrSA",
+            "7I4Y2HtAwlY",
+            "us-andJDfWw",
+            "Q0Bzn1teWfo",
+            "pADxBfOrYeQ",
+            "KnA1tIR0qCE",
+            "lBNPfBcb6pw",
+            "zt9MIiBfxF8",
+            "FwX12hPqCws",
+            "6e2LXGrze2M",
+            "0dvclhcFjSo",
+            "UTkuDYOIxKo",
+            "4n044cOTiNg",
+            "VxksG15OgRY",
+            "4193MkWdJcQ",
+            "2bDX4Uzy8Wc",
+            "SkyBNX2auEc",
+            "JrSLspuJ2gk",
+            "GZO7V-CWSJ8",
+            "TYMIy4qlRzc",
+            "jK7q_Li7LU0",
+            "UiTIXDQEkPM"
+        ],
+        "test_session_set": [
+            "DTGtC7FC4oI",
+            "CIWvk3gJ_7E",
+            "A1fZ6rlOwxk",
+            "cB8WCrmJEEA",
+            "gGQ_2Ank9-U",
+            "b91abcyMeNY",
+            "9MuaFeyczM0",
+            "rwZxhdK_UVU",
+            "hd_WbrqYnTE",
+            "4BWcqA9cFFk",
+            "MuOkgCkob6E",
+            "OdhFWA3Fgb0",
+            "Pe3ZabrUiUk",
+            "U0c_uv8DU08",
+            "kTqt_3YJBbs",
+            "eWXVVMi4bf8",
+            "RCD-vXumMIs",
+            "DuwiezJCZ3I",
+            "BNgEqjAiKU4",
+            "Cbn0t7eqaEs",
+            "D3M9Q0WwuFY",
+            "oO-Bv4iRZKI",
+            "bBvqtbZdSZQ",
+            "-mnoItCi32Y",
+            "yQPxKfz-bwI",
+            "ITVZ4cx5big",
+            "uY5_e0L7mjw",
+            "J_Pl9ftnNXE",
+            "Jt_pb40MRy4",
+            "E3hAj2niRNo",
+            "kl2xVSxVECQ",
+            "gzebuFs90Xg",
+            "6DQSIzJ9J-g",
+            "pJJlnryTYWg",
+            "X87qsfVLFoc",
+            "W7BIPWU1WHc",
+            "Oj506EVOiQc",
+            "DPEvWvcstxA",
+            "Pk16UT5-zeM",
+            "XMK1KEk5OdU",
+            "PpNm5llsEW8",
+            "sAp3rvablGk"
+        ]
+    }
+}

third_party/EgoSpeak/MiniROAD/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .dataset_builder import build_dataset, build_data_loader
+from .dataset import THUMOSDataset, FINEACTIONDataset

third_party/EgoSpeak/MiniROAD/datasets/dataset.py

@@ -0,0 +1,314 @@
+import gc
+import json
+import os.path as osp
+import sys
+
+import numpy as np
+import torch
+import torch.utils.data as data
+from configs.feature_sizes import FEATURE_SIZES
+from datasets.dataset_builder import DATA_LAYERS
+from utils.util import parse_flow_type, parse_rgb_type
+
+
+@DATA_LAYERS.register("THUMOS")
+@DATA_LAYERS.register("TVSERIES")
+@DATA_LAYERS.register("EasyCom")
+@DATA_LAYERS.register("EGO4D")
+@DATA_LAYERS.register("IEMOCAP")
+@DATA_LAYERS.register("IEMOCAP_DEBUG")
+@DATA_LAYERS.register("YTConv")
+class THUMOSDataset(data.Dataset):
+    
+    def __init__(self, cfg, mode='train'):
+        self.root_path = cfg['root_path']
+        self.mode = mode
+        self.training = mode == 'train'
+        self.window_size = cfg['window_size']
+        self.enc_steps = cfg['window_size'] if 'window_size' in cfg else None
+        self.dec_steps = cfg['query_num'] if 'query_num' in cfg else None
+        
+        self.stride = cfg['stride']
+        data_name = cfg['data_name']
+        self.model = cfg['model']
+        self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set']                      
+        self.num_classes = cfg['num_classes']
+        self.inputs = []
+
+        self._load_features(cfg)
+        self._init_features()
+        
+    def _load_features(self, cfg):
+        self.annotation_type = cfg['annotation_type']
+        self.rgb_type = cfg['rgb_type']
+        self.flow_type = cfg['flow_type']
+        self.target_all = {}
+        self.rgb_inputs = {}
+        self.flow_inputs = {}
+        dummy_target = np.zeros((self.window_size-1, self.num_classes))
+        
+        rgb_type = parse_rgb_type(cfg['rgb_type'])
+        flow_type = parse_flow_type(cfg['flow_type'])
+        
+        dummy_rgb = np.zeros((self.window_size-1, FEATURE_SIZES[rgb_type]))
+        dummy_flow = np.zeros((self.window_size-1, FEATURE_SIZES[flow_type]))
+        for vid in self.vids:
+            target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy'))
+            rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy'))
+            flow = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy'))
+                                                  
+            if self.training:
+                self.target_all[vid] = np.concatenate((dummy_target, target), axis=0)
+                self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0)
+                self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0)
+            else:
+                self.target_all[vid] = target
+                self.rgb_inputs[vid] = rgb
+                self.flow_inputs[vid] = flow
+    
+    def _init_features(self):
+        del self.inputs
+        gc.collect()
+        self.inputs = []
+        if self.model == 'MiniROAD':
+            for vid in self.vids:
+                target = self.target_all[vid]
+                                                         
+                
+                if self.training:
+                    seed = np.random.randint(self.stride)
+                    for start, end in zip(range(seed, target.shape[0], self.stride), 
+                        range(seed + self.window_size, target.shape[0]+1, self.stride)):
+                        self.inputs.append([
+                            vid, start, end, target[start:end]
+                        ])
+
+                else:
+                    start = 0
+                    end = target.shape[0]
+                    self.inputs.append([
+                        vid, start, end, target[start:end]
+                    ])
+                    
+        elif self.model == 'OADTR':
+            for vid in self.vids:
+                target = self.target_all[vid]
+                                                   
+                                                                       
+                seed = np.random.randint(self.stride) if self.training else 0
+                                                   
+                for start, end in zip(
+                    range(seed, target.shape[0], 1), 
+                    range(seed + self.window_size, target.shape[0]+1, 1)
+                    ):
+                    enc_target = target[start:end]
+                    dec_target = target[end:end + self.dec_steps]
+                    distance_target, class_h_target = self.get_distance_target(target[start:end])
+                    
+                    self.inputs.append([
+                        vid, start, end, enc_target, distance_target, class_h_target, dec_target
+                    ])
+            
+        else:               
+            for vid in self.vids:
+                target = self.target_all[vid]
+                seed = np.random.randint(self.stride) if self.training else 0
+                for start, end in zip(
+                    range(seed, target.shape[0], 1), 
+                    range(seed + self.window_size, target.shape[0]+1, 1)
+                    ):
+                    self.inputs.append([
+                        vid, start, end, target[start:end]
+                    ])
+    
+    def get_distance_target(self, target_vector):
+        target_matrix = np.zeros(self.enc_steps - 1)
+        target_argmax = target_vector[self.enc_steps-1].argmax()
+        for i in range(self.enc_steps - 1):
+            if target_vector[i].argmax() == target_argmax:
+                target_matrix[i] = 1.
+        return target_matrix, target_vector[self.enc_steps-1]
+    
+    def __getitem__(self, index):
+        if self.model == 'MiniROAD' or self.model == 'Transformer':
+            vid, start, end, target = self.inputs[index]
+            rgb_input = self.rgb_inputs[vid][start:end]
+            flow_input = self.flow_inputs[vid][start:end]
+            rgb_input = torch.tensor(rgb_input.astype(np.float32))
+            flow_input = torch.tensor(flow_input.astype(np.float32))
+            target = torch.tensor(target.astype(np.float32))
+            return vid, rgb_input, flow_input, target
+        
+        elif self.model == 'OADTR':
+            vid, start, end, enc_target, distance_target, class_h_target, dec_target = self.inputs[index]
+            rgb_input = self.rgb_inputs[vid][start:end]
+            flow_input = self.flow_inputs[vid][start:end]
+            rgb_input = torch.tensor(rgb_input.astype(np.float32))
+            flow_input = torch.tensor(flow_input.astype(np.float32))
+            enc_target = torch.tensor(enc_target.astype(np.float32))
+            distance_target = torch.tensor(distance_target.astype(np.float32))
+            class_h_target = torch.tensor(class_h_target.astype(np.float32))
+            dec_target = torch.tensor(dec_target.astype(np.float32))
+            return vid, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target
+
+    def __len__(self):
+        return len(self.inputs)
+    
+
+@DATA_LAYERS.register("THUMOS_ANTICIPATION")
+@DATA_LAYERS.register("TVSERIES_ANTICIPATION")
+@DATA_LAYERS.register("IEMOCAP_ANTICIPATION")
+@DATA_LAYERS.register("EasyCom_ANTICIPATION")
+@DATA_LAYERS.register("YTConv_ANTICIPATION")
+@DATA_LAYERS.register("EGO4D_ANTICIPATION")
+class THUMOSDataset(data.Dataset):
+    
+    def __init__(self, cfg, mode='train'):
+        self.cfg = cfg
+        self.root_path = cfg['root_path']
+        self.mode = mode
+        self.training = mode == 'train'
+        self.window_size = cfg['window_size']
+        self.stride = cfg['stride']
+        self.anticipation_length = cfg['anticipation_length']
+        data_name = cfg["data_name"].split('_')[0]
+        if 'data_info_name' in cfg.keys() and cfg['data_info_name'] is not None:
+            data_name = cfg['data_info_name']
+        
+        self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set']                      
+        print(f'data_name: {data_name}')
+        print(f'{mode} session set length: {len(self.vids)}')
+        self.num_classes = cfg['num_classes']
+        self.inputs = []
+        
+        self._load_features(cfg)
+        self._init_features()
+        
+    def _load_features(self, cfg):
+        self.annotation_type = cfg['annotation_type']
+        self.rgb_type = cfg['rgb_type']
+        self.flow_type = cfg['flow_type']
+        
+        self.target_all = {}
+        self.rgb_inputs = {}
+        self.flow_inputs = {}
+        
+        dummy_target = np.zeros((self.window_size-1, self.num_classes))
+        dummy_rgb = np.zeros((self.window_size-1, FEATURE_SIZES[cfg['rgb_type']]))
+        dummy_flow = np.zeros((self.window_size-1, FEATURE_SIZES[cfg['flow_type']]))
+        
+        for vid in self.vids:
+            target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy'))
+            rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy'))
+            flow = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy'))
+            
+            if self.training:
+                self.target_all[vid] = np.concatenate((dummy_target, target), axis=0)
+                self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0)
+                self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0)
+            else:                      
+                self.target_all[vid] = target
+                self.rgb_inputs[vid] = rgb
+                self.flow_inputs[vid] = flow
+        
+    def _init_features(self):
+        del self.inputs
+        gc.collect()
+        self.inputs = []
+
+        for vid in self.vids:
+            target = self.target_all[vid]
+            if self.training:
+                seed = np.random.randint(self.stride)
+                for start, end in zip(range(seed, target.shape[0], self.stride), 
+                    range(seed + self.window_size, target.shape[0]-self.anticipation_length, self.stride)):
+                    self.inputs.append([
+                        vid, start, end, target[start:end], target[end:end+self.anticipation_length]
+                    ])
+            else:
+                start = 0
+                end = target.shape[0] - self.anticipation_length
+                ant_target = []                                                                     
+                                                                                                        
+                for s in range(1, target.shape[0] - self.anticipation_length + 1):
+                    ant_target.append(target[s:s + self.anticipation_length])
+                
+                self.inputs.append([
+                    vid, start, end, target[start:end], np.array(ant_target)
+                ])
+                
+                                                
+    
+    def __getitem__(self, index):
+        vid, start, end, target, ant_target = self.inputs[index]
+                                                  
+        
+        rgb_input = self.rgb_inputs[vid][start:end]
+        flow_input = self.flow_inputs[vid][start:end]
+        
+        rgb_input = torch.tensor(rgb_input.astype(np.float32))
+        flow_input = torch.tensor(flow_input.astype(np.float32))
+        
+        target = torch.tensor(target.astype(np.float32))
+        ant_target = torch.tensor(ant_target.astype(np.float32))
+        
+        return vid, rgb_input, flow_input, target, ant_target
+
+    def __len__(self):
+        return len(self.inputs)                          
+    
+    
+@DATA_LAYERS.register("FINEACTION")
+class FINEACTIONDataset(data.Dataset):
+    
+    def __init__(self, cfg, mode='train'):
+        self.root_path = cfg['root_path']
+        self.mode = mode
+        self.training = mode == 'train'
+        self.window_size = cfg['window_size']
+        self.stride = cfg['stride']
+        data_name = cfg['data_name']
+        self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set']                      
+        self.num_classes = cfg['num_classes']
+        self.inputs = []
+        self._load_features(cfg)
+        self._init_features()
+
+    def _load_features(self, cfg):
+        self.annotation_type = cfg['annotation_type']
+        self.rgb_type = cfg['rgb_type']
+        self.flow_type = cfg['flow_type']
+        
+    def _init_features(self, seed=0):
+                          
+        del self.inputs
+        gc.collect()
+        self.inputs = []
+        for vid in self.vids:
+            target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy'))
+            if self.training:
+                seed = np.random.randint(self.stride)
+                for start, end in zip(range(seed, target.shape[0], self.stride), 
+                    range(seed + self.window_size, target.shape[0]+1, self.stride)):
+                    self.inputs.append([
+                        vid, start, end
+                    ])
+            else:
+                start = 0
+                end = target.shape[0]
+                self.inputs.append([
+                    vid, start, end
+                ])
+
+    def __getitem__(self, index):
+        vid, start, end = self.inputs[index]
+        rgb_input = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy'), mmap_mode='r')[start:end]
+        flow_input = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy'), mmap_mode='r')[start:end]
+        target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy'), mmap_mode='r')[start:end]
+        rgb_input = torch.tensor(rgb_input.astype(np.float32))
+        flow_input = torch.tensor(flow_input.astype(np.float32))
+        target = torch.tensor(target.astype(np.float32))
+        return rgb_input, flow_input, target
+
+    def __len__(self):
+        return len(self.inputs)    

third_party/EgoSpeak/MiniROAD/datasets/dataset_builder.py

@@ -0,0 +1,32 @@
+__all__ = [
+    'build_dataset',
+    'build_data_loader',
+]
+
+import torch.utils.data as data
+from utils import Registry
+
+DATA_LAYERS = Registry()
+
+def build_dataset(cfg):
+    data_layer = DATA_LAYERS[f'{cfg["data_name"]}']
+    return data_layer
+
+def build_data_loader(cfg, mode):
+    data_layer = build_dataset(cfg)
+    data_loader = data.DataLoader(
+        dataset=data_layer(cfg, mode),
+        batch_size=cfg["batch_size"] if mode == 'train' else cfg["test_batch_size"],
+        shuffle=True if mode == 'train' else False,
+        num_workers=cfg["num_workers"],
+        pin_memory=False,                                                  
+    )
+                                                                  
+    
+                          
+                                                                                                          
+                                                                          
+                                                                                                       
+                                
+                                                                        
+    return data_loader

third_party/EgoSpeak/MiniROAD/main.py

@@ -0,0 +1,294 @@
+import argparse
+import os
+import os.path as osp
+
+import torch
+import wandb
+import yaml
+from criterions import build_criterion
+from model import build_model
+from torch.utils.tensorboard import SummaryWriter
+from trainer import build_eval, build_trainer
+from utils import *
+from utils import get_logger
+
+from datasets import build_data_loader
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str)
+    parser.add_argument("--eval", action="store_true")
+    parser.add_argument("--ckpt", type=str)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument("--tensorboard", action="store_true")
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--lr_scheduler", action="store_true")
+    parser.add_argument("--no_rgb", action="store_true")
+    parser.add_argument("--no_flow", action="store_true")
+    parser.add_argument(
+        "--save-as-cfg",
+        action="store_true",
+        help="Save outputs with config name in the output directory",
+    )
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument(
+        "--save-preds",
+        action="store_true",
+        help="Save predictions in the output directory",
+    )
+    # parser.add_argument('--model', type=str, default=None)
+    args = parser.parse_args()
+
+    # combine argparse and yaml
+    opt = yaml.load(open(args.config), Loader=yaml.FullLoader)
+    opt.update(vars(args))
+    cfg = opt
+
+    if "checkpoint" in cfg and cfg["checkpoint"] is not None and args.ckpt is None:
+        args.ckpt = cfg["checkpoint"]
+        print(f"checkpoint is given in the config, set args.ckpt to {args.ckpt}")
+
+    seed = 20 if "seed" not in cfg else cfg["seed"]
+    set_seed(seed)
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    annotation_type = os.path.basename(cfg["annotation_type"])
+
+    identifier = os.path.basename(cfg["config"])[:-5]
+    identifier += "_"
+    if cfg["no_flow"] == False:
+        identifier += "A"
+    if cfg["no_rgb"] == False:
+        identifier += "V"
+
+    if cfg["save_as_cfg"] == True:
+        output_dir = osp.join("output_cfg", identifier)
+        result_path = create_outdir(output_dir, overwrite=True)
+    else:
+        identifier = get_identifier(cfg)
+        annotation_type = os.path.basename(cfg["annotation_type"])
+        identifier += f"_{annotation_type}"
+        result_path = create_outdir(osp.join(cfg["output_path"], identifier))
+
+    logger = get_logger(result_path)
+    logger.info(f"Identifier: {identifier}")
+    logger.info(f"result_path: {result_path}")
+    logger.info(cfg)
+
+    testloader = build_data_loader(cfg, mode="test")
+    print("testloader loaded")
+    model = build_model(cfg, device)
+    evaluate = build_eval(cfg)
+
+    trainloader = build_data_loader(cfg, mode="train")
+    print("trainloader loaded")
+    
+    # evaluation
+    if args.eval == True and cfg["ckpt"] == None:
+        print("Automatically find the best checkpoint")
+        ckpt_root = os.path.join("output_cfg", identifier, "ckpts")
+        ckpt_dirs = os.listdir(ckpt_root)
+
+        ckpt_path = os.path.join(ckpt_root, "best.pth")
+        if os.path.exists(ckpt_path):
+            args.ckpt = ckpt_path
+            print(f"args.ckpt set to {args.ckpt}")
+
+        else:
+            ckpt_path = os.path.join(ckpt_root, "pred", "best.pth")
+            if os.path.exists(ckpt_path):
+                args.ckpt = ckpt_path
+                print(f"args.ckpt set to {args.ckpt}")
+
+        print("Warninig) no ckpt_path found, there can be multiple ckpts")
+        print(f"ckpt_dirs: {ckpt_dirs}")
+        pred_dir_number = 0
+        ckpt_dirs.sort(reverse=True)
+
+        if args.ckpt == None:
+            for ckpt_dir in ckpt_dirs:
+                ckpt_path = os.path.join(ckpt_root, ckpt_dir, "best.pth")
+
+                if os.path.exists(ckpt_path):
+                    args.ckpt = ckpt_path
+                    print(f"Found ckpt, args.ckpt was None, set to {args.ckpt}")
+                    break
+
+        if args.ckpt == None:
+            raise ValueError("No checkpoint found, args.ckpt is None")
+
+    # load checkpoint if args.ckpt is given
+    if args.eval == True:
+        if args.ckpt == None:
+            raise ValueError("Please specify the checkpoint path")
+        checkpoint = torch.load(args.ckpt)
+        model.load_state_dict(checkpoint)
+
+    elif args.eval == False and args.ckpt != None:
+        checkpoint = torch.load(args.ckpt)
+        print(f"Load checkpoint from {args.ckpt}")
+        model_dict = model.state_dict()
+
+        unnecessary_keys = []
+        matched_ckpt = {}
+        # Filter out unnecessary keys
+        for k, v in checkpoint.items():
+            if not (k in model_dict and model_dict[k].shape == checkpoint[k].shape):
+                unnecessary_keys.append(k)
+            elif (
+                "except_classifier" in cfg
+                and cfg["except_classifier"] == True
+                and "classification" in k
+            ):
+                print(f"Unnecessary key: {k} is classifier")
+                unnecessary_keys.append(k)
+            else:
+                matched_ckpt[k] = v
+
+        # Overwrite entries in the existing state dict
+        model_dict.update(matched_ckpt)
+
+        # Load the new state dict
+        model.load_state_dict(model_dict)
+
+        print(f"Unnecessary keys: {unnecessary_keys}")
+        
+    if cfg["wandb"] == True:
+        # args.tensorboard = True
+        print("wandb init")
+        os.environ["WANDB__SERVICE_WAIT"] = "3000"
+        wandb.init(
+            project="egospeak",
+            tags=[cfg["model"], cfg["data_name"], annotation_type],
+            config=cfg,
+            name=os.path.basename(result_path),
+        )
+
+    if args.eval == True or args.debug == True:
+        if args.eval == True and args.ckpt == None:
+            raise ValueError("Please specify the checkpoint path")
+
+        logger.info(f'Dataset: {cfg["data_name"]},  Model: {cfg["model"]}')
+        logger.info(f"Result path:{result_path}")
+
+        for phase in ["test"]:
+            if phase == "train":
+                raise NotImplementedError("evaluation on train split is not implemented")
+                dataloader = trainloader
+                
+            elif phase == "test":
+                dataloader = testloader
+
+            mAP, mcAP, preds = evaluate(
+                model, dataloader, logger, result_path=result_path
+            )
+
+            pred_path = osp.join(result_path, "eval")
+            os.makedirs(pred_path, exist_ok=True)
+
+            if args.save_preds == True:
+                print(f"pred_path: {pred_path}")
+                for k, v in preds.items():
+                    np.save(osp.join(pred_path, f"{k}.npy"), v)
+
+            logger.info(f'{phase}) {cfg["task"]} result: {mAP*100:.2f} mAP')
+            logger.info(f'{phase}) {cfg["task"]} result: {mcAP*100:.2f} mcAP')
+        exit()
+
+    criterion = build_criterion(cfg, device)
+    train_one_epoch = build_trainer(cfg)
+    optim = torch.optim.AdamW if cfg["optimizer"] == "AdamW" else torch.optim.Adam
+    optimizer = optim(
+        [{"params": model.parameters(), "initial_lr": cfg["lr"]}],
+        lr=cfg["lr"],
+        weight_decay=cfg["weight_decay"],
+    )
+
+    scheduler = (
+        build_lr_scheduler(cfg, optimizer, len(trainloader))
+        if args.lr_scheduler
+        else None
+    )
+    writer = SummaryWriter(osp.join(result_path, "runs")) if args.tensorboard else None
+    scaler = torch.cuda.amp.GradScaler() if args.amp else None
+    total_params = sum(p.numel() for p in model.parameters())
+
+    logger.info(f'Dataset: {cfg["data_name"]},  Model: {cfg["model"]}')
+    logger.info(
+        f'lr:{cfg["lr"]} | Weight Decay:{cfg["weight_decay"]} | Window Size:{cfg["window_size"]} | Batch Size:{cfg["batch_size"]}'
+    )
+    logger.info(
+        f'Total epoch:{cfg["num_epoch"]} | Total Params:{total_params/1e6:.1f} M | Optimizer: {cfg["optimizer"]}'
+    )
+    logger.info(f"Output Path:{result_path}")
+
+    best_mAP, best_epoch = 0, 0
+    best_cAP, best_cAP_epoch = 0, 0
+
+    pred_save_path = osp.join(result_path, "ckpts", "pred")
+    print(f"pred_save_path: {pred_save_path}")
+
+    # if you want to overwrite, change the code
+    os.makedirs(pred_save_path, exist_ok=True)
+
+    for epoch in range(1, cfg["num_epoch"] + 1):
+        epoch_loss = train_one_epoch(
+            trainloader,
+            model,
+            criterion,
+            optimizer,
+            scaler,
+            epoch,
+            cfg["wandb"],
+            writer,
+            scheduler=scheduler,
+        )
+        trainloader.dataset._init_features()
+
+        mAP, cAP, preds = evaluate(
+            model, testloader, logger, criterion, cfg["wandb"], epoch, writer
+        )
+        
+        if mAP > best_mAP:
+            best_mAP = mAP
+            best_epoch = epoch
+            torch.save(model.state_dict(), osp.join(pred_save_path, "best.pth"))
+            
+            if args.save_preds == True:
+                for k, v in preds.items():
+                    np.save(osp.join(pred_save_path, f"{k}.npy"), v)
+
+        if cAP > best_cAP:
+            best_cAP = cAP
+            best_cAP_epoch = epoch
+            torch.save(model.state_dict(), osp.join(pred_save_path, "best_cAP.pth"))
+
+        logger.info(
+            f'Epoch {epoch} mAP: {mAP*100:.2f} | Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}'
+        )
+        logger.info(
+            f'Epoch {epoch} cAP: {cAP*100:.2f} | Best cAP: {best_cAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}'
+        )
+
+
+        if cfg["wandb"] == True:
+            logger.info("Wandb logging")
+            wandb.log(
+                {
+                    "Epoch": epoch,
+                    "Train Loss (Epoch)": epoch_loss / len(trainloader),
+                    "Learning Rate": optimizer.param_groups[0]["lr"],
+                    "anticipation average mAP": mAP,
+                    "anticipation average cAP": cAP,
+                }
+            )
+
+    # os.rename(osp.join(result_path, 'ckpts', 'best.pth'), osp.join(result_path, 'ckpts', f'best_{best_mAP*100:.2f}.pth'))
+    # save the training result and config as txt file
+    with open(osp.join(result_path, "result.txt"), "w") as f:
+        f.write(f"Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}\n")
+        f.write(f"Best cAP: {best_cAP*100:.2f} at epoch {best_cAP_epoch}\n")
+        f.write(f"Config: {cfg}")
+
+    wandb.finish() if args.wandb else None

third_party/EgoSpeak/MiniROAD/model/__init__.py

@@ -0,0 +1,4 @@
+from .mamba import OADMamba
+from .model_builder import build_model
+from .rnn import MROAD
+from .transformer_models import ViTEnc

third_party/EgoSpeak/MiniROAD/model/mamba/__init__.py

@@ -0,0 +1 @@
+from .mamba import OADMamba, OADMambaA

third_party/EgoSpeak/MiniROAD/model/mamba/mamba.py

@@ -0,0 +1,172 @@
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from mamba_ssm import Mamba
+# from mambapy.mamba import Mamba, MambaConfig
+from model.model_builder import META_ARCHITECTURES
+from utils.util import parse_flow_type, parse_rgb_type
+
+
+@META_ARCHITECTURES.register("Mamba")
+class OADMamba(nn.Module):
+    
+    def __init__(self, cfg):
+        super(OADMamba, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.input_dim = 0
+        
+        rgb_type = parse_rgb_type(cfg['rgb_type'])
+        flow_type = parse_flow_type(cfg['flow_type'])
+        
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[rgb_type]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[flow_type]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        
+        # raise NotImplementedError("The following code is not implemented yet")
+        self.out_dim = cfg['num_classes']
+        self.window_size = cfg['window_size']
+        
+        self.relu = nn.ReLU()
+        
+        self.embedding_dim = cfg['embedding_dim']
+        self.d_state = cfg['d_state']
+        self.d_conv = cfg['d_conv']
+        self.expand = cfg['expand']
+
+        # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers)
+        # self.mamba = Mamba(self.mamba_config)
+        self.mamba = Mamba(
+            d_model=self.embedding_dim, # Model dimension d_model
+            d_state=self.d_state,  # SSM state expansion factor
+            d_conv=self.d_conv,    # Local convolution width
+            expand=self.expand,    # Block expansion factor
+        )
+        
+        # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout']),
+        )
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim)
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            x = torch.cat((x, flow_input), dim=2)
+        
+        x = self.layer1(x)
+        ht = self.mamba(x)
+        
+        logits = self.f_classification(ht)
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            out_dict['logits'] = pred_scores
+        return out_dict
+
+@META_ARCHITECTURES.register("MambaA")
+class OADMambaA(nn.Module):
+    
+    def __init__(self, cfg):
+        super(OADMambaA, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.input_dim = 0
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[cfg['flow_type']]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.anticipation_length = cfg["anticipation_length"]
+        self.out_dim = cfg['num_classes']
+        
+        self.embedding_dim = cfg['embedding_dim']
+        self.d_state = cfg['d_state'] if 'd_state' in cfg else 16
+        self.d_conv = cfg['d_conv'] if 'd_conv' in cfg else 4
+        self.expand = cfg['expand'] if 'expand' in cfg else 2
+        
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout'])
+        )
+        self.actionness = cfg['actionness']
+        if self.actionness:
+            self.f_actionness = nn.Sequential(
+                nn.Linear(self.hidden_dim, 1),
+            )
+        self.relu = nn.ReLU()
+        # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers)
+        # self.mamba = Mamba(self.mamba_config)
+        self.mamba = Mamba(
+            d_model=self.embedding_dim, # Model dimension d_model
+            d_state=self.d_state,  # SSM state expansion factor
+            d_conv=self.d_conv,    # Local convolution width
+            expand=self.expand,    # Block expansion facto
+        )
+        
+        # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.embedding_dim, self.out_dim)
+        )
+        self.anticipation_layer = nn.Sequential(
+            nn.Linear(self.embedding_dim, self.anticipation_length*self.embedding_dim),
+        )
+        self.cfg = cfg
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            # [64, 128, 2048]
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            # [64, 128, 5120]
+            x = torch.cat((x, flow_input), dim=2)
+
+        # audio channel is too huge to directly concat...
+        # to c
+
+        # B, L, D: mamba
+        B, S, feat_dim = x.shape # batch, length, dim
+        # 
+        x = self.layer1(x)
+        ht = self.mamba(x)
+        # h0 = torch.zeros(1, B, self.hidden_dim).to(x.device)
+        # ht, _ = self.gru(x, h0)
+        logits = self.f_classification(self.relu(ht))
+        anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.embedding_dim)
+        anticipation_logits = self.f_classification(self.relu(anticipation_ht))
+        
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+            out_dict['anticipation_logits'] = anticipation_logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1)
+            out_dict['logits'] = pred_scores
+            out_dict['anticipation_logits'] = pred_anticipation_scores
+
+        return out_dict

third_party/EgoSpeak/MiniROAD/model/model_builder.py

@@ -0,0 +1,9 @@
+__all__ = ['build_model']
+
+from utils import Registry
+
+META_ARCHITECTURES = Registry()
+
+def build_model(cfg, device=None):
+    model = META_ARCHITECTURES[cfg["model"]](cfg)
+    return model.to(device)

third_party/EgoSpeak/MiniROAD/model/rnn/__init__.py

@@ -0,0 +1 @@
+from .rnn import MROAD, MROADA

third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py

@@ -0,0 +1,153 @@
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from model.model_builder import META_ARCHITECTURES
+from utils.util import parse_flow_type, parse_rgb_type
+
+
+@META_ARCHITECTURES.register("MiniROAD")
+class MROAD(nn.Module):
+    
+    def __init__(self, cfg):
+        super(MROAD, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.input_dim = 0
+        
+        rgb_type = parse_rgb_type(cfg['rgb_type'])
+        flow_type = parse_flow_type(cfg['flow_type'])
+        
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[rgb_type]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[flow_type]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.out_dim = cfg['num_classes']
+        self.window_size = cfg['window_size']
+
+        self.relu = nn.ReLU()
+        self.embedding_dim = cfg['embedding_dim']
+        self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout']),
+        )
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim)
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            x = torch.cat((x, flow_input), dim=2)
+
+        # if self.use_rgb and self.use_flow:
+        #     x = torch.cat((rgb_input, flow_input), 2)
+        # elif self.use_rgb:
+        #     x = rgb_input
+        # elif self.use_flow:
+        #     x = flow_input
+        
+        x = self.layer1(x)
+        B, _, _ = x.shape
+        h0 = self.h0.expand(-1, B, -1).to(x.device)
+        ht, _ = self.gru(x, h0)
+        ht = self.relu(ht)
+        # ht = self.relu(ht + x)
+        logits = self.f_classification(ht)
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            out_dict['logits'] = pred_scores
+        return out_dict
+
+@META_ARCHITECTURES.register("MiniROADA")
+class MROADA(nn.Module):
+    
+    def __init__(self, cfg):
+        super(MROADA, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.input_dim = 0
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[cfg['flow_type']]
+
+        self.embedding_dim = cfg['embedding_dim']
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.anticipation_length = cfg["anticipation_length"]
+        self.out_dim = cfg['num_classes']
+
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout'])
+        )
+        
+        self.actionness = cfg['actionness']
+        if self.actionness:
+            self.f_actionness = nn.Sequential(
+                nn.Linear(self.hidden_dim, 1),
+            )
+        self.relu = nn.ReLU()
+        self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        self.anticipation_layer = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.anticipation_length*self.hidden_dim),
+        )
+        self.cfg = cfg
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            # [64, 128, 2048]
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            # [64, 128, 5120]
+            x = torch.cat((x, flow_input), dim=2)
+
+        # audio channel is too huge to directly concat...
+        # to c
+
+        B, S, feat_dim = x.shape # batch, length, dim
+        # 
+        x = self.layer1(x)
+        h0 = torch.zeros(1, B, self.hidden_dim).to(x.device)
+        ht, _ = self.gru(x, h0)
+        ht = self.relu(ht)
+        
+        logits = self.f_classification(self.relu(ht))
+        anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.hidden_dim)
+        anticipation_logits = self.f_classification(self.relu(anticipation_ht))
+        
+        
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+            out_dict['anticipation_logits'] = anticipation_logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1)
+            out_dict['logits'] = pred_scores
+            out_dict['anticipation_logits'] = pred_anticipation_scores
+
+        return out_dict

third_party/EgoSpeak/MiniROAD/model/transformer_models/Attention.py

@@ -0,0 +1,155 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0
+    ):
+        super().__init__()
+        self.num_heads = heads
+        head_dim = dim // heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(dropout_rate)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class AxialAttention(nn.Module):
+    def __init__(
+        self,
+        in_planes,
+        out_planes,
+        groups=8,
+        kernel_size=56,
+        stride=1,
+        bias=False,
+        width=False,
+    ):
+        assert (in_planes % groups == 0) and (out_planes % groups == 0)
+        super(AxialAttention, self).__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.groups = groups
+        self.group_planes = out_planes // groups
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.bias = bias
+        self.width = width
+
+        # Multi-head self attention
+        self.qkv_transform = nn.Conv1d(
+            in_planes,
+            out_planes * 2,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
+        self.bn_similarity = nn.BatchNorm2d(groups * 3)
+        self.bn_output = nn.BatchNorm1d(out_planes * 2)
+
+        # Position embedding
+        self.relative = nn.Parameter(
+            torch.randn(self.group_planes * 2, kernel_size * 2 - 1),
+            requires_grad=True,
+        )
+        query_index = torch.arange(kernel_size).unsqueeze(0)
+        key_index = torch.arange(kernel_size).unsqueeze(1)
+        relative_index = key_index - query_index + kernel_size - 1
+        self.register_buffer('flatten_index', relative_index.view(-1))
+        if stride > 1:
+            self.pooling = nn.AvgPool2d(stride, stride=stride)
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        if self.width:
+            x = x.permute(0, 2, 1, 3)
+        else:
+            x = x.permute(0, 3, 1, 2)  # N, W, C, H
+        N, W, C, H = x.shape
+        x = x.contiguous().view(N * W, C, H)
+
+        # Transformations
+        qkv = self.bn_qkv(self.qkv_transform(x))
+        q, k, v = torch.split(
+            qkv.reshape(N * W, self.groups, self.group_planes * 2, H),
+            [self.group_planes // 2, self.group_planes // 2, self.group_planes],
+            dim=2,
+        )
+
+        # Calculate position embedding
+        all_embeddings = torch.index_select(
+            self.relative, 1, self.flatten_index
+        ).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
+        q_embedding, k_embedding, v_embedding = torch.split(
+            all_embeddings,
+            [self.group_planes // 2, self.group_planes // 2, self.group_planes],
+            dim=0,
+        )
+        qr = torch.einsum('bgci,cij->bgij', q, q_embedding)
+        kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3)
+        qk = torch.einsum('bgci, bgcj->bgij', q, k)
+        stacked_similarity = torch.cat([qk, qr, kr], dim=1)
+        stacked_similarity = (
+            self.bn_similarity(stacked_similarity)
+            .view(N * W, 3, self.groups, H, H)
+            .sum(dim=1)
+        )
+
+        # (N, groups, H, H, W)
+        similarity = F.softmax(stacked_similarity, dim=3)
+        sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
+        sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding)
+        stacked_output = torch.cat([sv, sve], dim=-1).view(
+            N * W, self.out_planes * 2, H
+        )
+        output = (
+            self.bn_output(stacked_output)
+            .view(N, W, self.out_planes, 2, H)
+            .sum(dim=-2)
+        )
+
+        if self.width:
+            output = output.permute(0, 2, 1, 3)
+        else:
+            output = output.permute(0, 2, 3, 1)
+
+        if self.stride > 1:
+            output = self.pooling(output)
+
+        return output
+
+    def reset_parameters(self):
+        self.qkv_transform.weight.data.normal_(
+            0, math.sqrt(1.0 / self.in_planes)
+        )
+        nn.init.normal_(self.relative, 0.0, math.sqrt(1.0 / self.group_planes))

third_party/EgoSpeak/MiniROAD/model/transformer_models/AxialNet.py

@@ -0,0 +1,244 @@
+import torch
+import torch.nn as nn
+from .Attention import AxialAttention
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class AxialBlock(nn.Module):
+    expansion = 2
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        downsample=None,
+        groups=1,
+        base_width=64,
+        dilation=1,
+        norm_layer=None,
+        kernel_size=56,
+    ):
+        super(AxialBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.0))
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv_down = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.hight_block = AxialAttention(
+            width, width, groups=groups, kernel_size=kernel_size
+        )
+        self.width_block = AxialAttention(
+            width,
+            width,
+            groups=groups,
+            kernel_size=kernel_size,
+            stride=stride,
+            width=True,
+        )
+        self.conv_up = conv1x1(width, planes * self.expansion)
+        self.bn2 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv_down(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.hight_block(out)
+        out = self.width_block(out)
+        out = self.relu(out)
+
+        out = self.conv_up(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class AxialAttentionNet(nn.Module):
+    def __init__(
+        self,
+        layers,
+        num_classes=1000,
+        zero_init_residual=True,
+        groups=8,
+        width_per_group=64,
+        replace_stride_with_dilation=None,
+        norm_layer=None,
+        s=0.5,
+    ):
+        super(AxialAttentionNet, self).__init__()
+        block = AxialBlock
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = int(64 * s)
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                "or a 3-element tuple, got {}".format(
+                    replace_stride_with_dilation
+                )
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(
+            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(
+            block, int(128 * s), layers[0], kernel_size=56
+        )
+        self.layer2 = self._make_layer(
+            block,
+            int(256 * s),
+            layers[1],
+            stride=2,
+            kernel_size=56,
+            dilate=replace_stride_with_dilation[0],
+        )
+        self.layer3 = self._make_layer(
+            block,
+            int(512 * s),
+            layers[2],
+            stride=2,
+            kernel_size=28,
+            dilate=replace_stride_with_dilation[1],
+        )
+        self.layer4 = self._make_layer(
+            block,
+            int(1024 * s),
+            layers[3],
+            stride=2,
+            kernel_size=14,
+            dilate=replace_stride_with_dilation[2],
+        )
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(int(1024 * block.expansion * s), num_classes)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+                if isinstance(m, nn.Conv1d):
+                    pass
+                else:
+                    nn.init.kaiming_normal_(
+                        m.weight, mode='fan_out', nonlinearity='relu'
+                    )
+            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, AxialBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(
+        self, block, planes, blocks, kernel_size=56, stride=1, dilate=False
+    ):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                stride,
+                downsample,
+                groups=self.groups,
+                base_width=self.base_width,
+                dilation=previous_dilation,
+                norm_layer=norm_layer,
+                kernel_size=kernel_size,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        if stride != 1:
+            kernel_size = kernel_size // 2
+
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                    kernel_size=kernel_size,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x, include_conv5=False, include_top=False):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        if include_conv5:
+            x = self.layer4(x)
+
+        if include_top:
+            x = self.avgpool(x)
+            x = torch.flatten(x, 1)
+            x = self.fc(x)
+
+        return x
+
+    def forward(self, x, include_conv5=False, include_top=False):
+        return self._forward_impl(x, include_conv5, include_top)
+
+
+if __name__ == "__main__":
+    model = AxialAttentionNet([1, 2, 4, 1], s=0.5)
+    # model = AxialAttentionNet([3, 4, 6, 3], s=0.25)
+    # print(model)
+
+    x = torch.randn([8, 3, 224, 224])
+    print(model(x).shape)
+    print(model(x, include_conv5=True).shape)
+    print(model(x, include_conv5=True, include_top=True).shape)

third_party/EgoSpeak/MiniROAD/model/transformer_models/BiT.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+
+class StdConv2d(nn.Conv2d):
+    def forward(self, x):
+        w = self.weight
+        v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False)
+        w = (w - m) / torch.sqrt(v + 1e-10)
+        return F.conv2d(
+            x,
+            w,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return StdConv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return StdConv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class PreActBottleneck(nn.Module):
+    """Pre-activation (v2) bottleneck block.
+    Follows the implementation of "Identity Mappings in Deep Residual Networks":
+    https://github.com/KaimingHe/resnet-1k-layers/blob/master/resnet-pre-act.lua
+    Except it puts the stride on 3x3 conv when available.
+    """
+
+    def __init__(self, in_planes, out_planes=None, mid_planes=None, stride=1):
+        super(PreActBottleneck, self).__init__()
+        out_planes = out_planes or in_planes
+        mid_planes = mid_planes or out_planes // 4
+
+        self.gn1 = nn.GroupNorm(32, in_planes)
+        self.conv1 = conv1x1(in_planes, mid_planes)
+        self.gn2 = nn.GroupNorm(32, mid_planes)
+        self.conv2 = conv3x3(mid_planes, mid_planes, stride)
+        self.gn3 = nn.GroupNorm(32, mid_planes)
+        self.conv3 = conv1x1(mid_planes, out_planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride != 1 or in_planes != out_planes:
+            # Projection also with pre-activation according to paper.
+            self.downsample = conv1x1(in_planes, out_planes, stride)
+
+    def forward(self, x):
+        out = self.relu(self.gn1(x))
+
+        # Residual branch
+        residual = x
+        if hasattr(self, 'downsample'):
+            residual = self.downsample(out)
+
+        # Unit's branch
+        out = self.conv1(out)
+        out = self.conv2(self.relu(self.gn2(out)))
+        out = self.conv3(self.relu(self.gn3(out)))
+
+        return out + residual
+
+
+class ResNetV2Model(nn.Module):
+    """Implementation of Pre-activation (v2) ResNet mode."""
+
+    def __init__(self, block_units, width_factor, head_size=21843):
+        super(ResNetV2Model, self).__init__()
+        wf = width_factor  # shortcut 'cause we'll use it a lot.
+
+        # The following will be unreadable if we split lines.
+        # pylint: disable=line-too-long
+        self.conv1 = nn.Sequential(OrderedDict([
+            ('conv', StdConv2d(3, 64*wf, kernel_size=7, stride=2, padding=3, bias=False)),
+            ('pad', nn.ConstantPad2d(1, 0)),
+            ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=0)),
+        ]))
+
+        self.conv2 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=64*wf, out_planes=256*wf, mid_planes=64*wf))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=256*wf, out_planes=256*wf, mid_planes=64*wf)) for i in range(2, block_units[0] + 1)],
+        ))
+        self.conv3 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=256*wf, out_planes=512*wf, mid_planes=128*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=512*wf, out_planes=512*wf, mid_planes=128*wf)) for i in range(2, block_units[1] + 1)],
+        ))
+        self.conv4 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=512*wf, out_planes=1024*wf, mid_planes=256*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=1024*wf, out_planes=1024*wf, mid_planes=256*wf)) for i in range(2, block_units[2] + 1)],
+        ))
+        self.conv5 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=1024*wf, out_planes=2048*wf, mid_planes=512*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=2048*wf, out_planes=2048*wf, mid_planes=512*wf)) for i in range(2, block_units[3] + 1)],
+        ))
+        # pylint: enable=line-too-long
+
+        self.head = nn.Sequential(OrderedDict([
+            ('gn', nn.GroupNorm(32, 2048*wf)),
+            ('relu', nn.ReLU(inplace=True)),
+            ('avg', nn.AdaptiveAvgPool2d(output_size=1)),
+            ('conv', nn.Conv2d(2048*wf, head_size, kernel_size=1, bias=True)),
+        ]))
+
+    def forward(self, x, include_conv5=False, include_top=False):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        if include_conv5:
+            x = self.conv5(x)
+        if include_top:
+            x = self.head(x)
+
+        if include_top and include_conv5:
+            assert x.shape[-2:] == (1,1,)
+            return x[..., 0, 0]
+
+        return x

third_party/EgoSpeak/MiniROAD/model/transformer_models/HybridViT.py

@@ -0,0 +1,223 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .BiT import ResNetV2Model
+from .AxialNet import AxialAttentionNet
+from .Transformer import TransformerModel
+from .PositionalEncoding import (
+    FixedPositionalEncoding,
+    LearnedPositionalEncoding,
+)
+
+
+class HybridVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5,
+        dropout_rate,
+        positional_encoding_type,
+        backbone=None,
+    ):
+        super(HybridVisionTransformer, self).__init__()
+
+        assert embedding_dim % num_heads == 0
+
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.out_dim = out_dim
+        self.num_channels = num_channels
+        self.include_conv5 = include_conv5
+        self.backbone = backbone
+
+        self.backbone_model, self.flatten_dim = self.configure_backbone()
+
+        self.projection_encoding = nn.Linear(self.flatten_dim, embedding_dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
+
+        self.decoder_dim = int(img_dim / 16.0) ** 2
+        if self.include_conv5:
+            self.decoder_dim = int(img_dim / 32.0) ** 2
+
+        self.decoder_dim += 1  # for the cls token
+
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.decoder_dim, self.embedding_dim, self.decoder_dim
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+
+        self.transformer = TransformerModel(
+            embedding_dim, num_layers, num_heads, hidden_dim
+        )
+        self.mlp_head = nn.Sequential(
+            nn.Linear(embedding_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+        self.to_cls_token = nn.Identity()
+
+    def forward(self, x):
+        # apply bit backbone
+        x = self.backbone_model(x, include_conv5=self.include_conv5)
+        x = x.view(x.size(0), -1, self.flatten_dim)
+
+        x = self.projection_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self.position_encoding(x)
+
+        # apply transformer
+        x = self.transformer(x)
+        x = self.to_cls_token(x[:, 0])
+        x = self.mlp_head(x)
+        x = F.log_softmax(x, dim=-1)
+
+        return x
+
+    def configure_backbone(self):
+        raise NotImplementedError("Method to be called in child class!!")
+
+
+class ResNetHybridViT(HybridVisionTransformer):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5=False,
+        dropout_rate=0.1,
+        positional_encoding_type="learned",
+        backbone='r50x1',
+    ):
+        super(ResNetHybridViT, self).__init__(
+            img_dim=img_dim,
+            out_dim=out_dim,
+            num_channels=num_channels,
+            embedding_dim=embedding_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            hidden_dim=hidden_dim,
+            include_conv5=include_conv5,
+            dropout_rate=dropout_rate,
+            positional_encoding_type=positional_encoding_type,
+            backbone=backbone,
+        )
+
+    def configure_backbone(self):
+        """
+        Current support offered for all BiT models
+        KNOWN_MODELS in https://github.com/google-research/big_transfer/blob/master/bit_pytorch/models.py
+
+        expects model name of style 'r{depth}x{width}'
+        where depth in [50, 101, 152]
+        where width in [1,2,3,4]
+        """
+        backbone = self.backbone
+        out_dim = self.out_dim
+
+        splits = backbone.split('x')
+        model_name = splits[0]
+        width_factor = int(splits[1])
+
+        if model_name in ['r50', 'r101'] and width_factor in [2, 4]:
+            return ValueError(
+                "Invalid Configuration of models -- expect 50x1, 50x3, 101x1, 101x3"
+            )
+        elif model_name == 'r152' and width_factor in [1, 3]:
+            return ValueError(
+                "Invalid Configuration of models -- expect 152x2, 152x4"
+            )
+
+        block_units_dict = {
+            'r50': [3, 4, 6, 3],
+            'r101': [3, 4, 23, 3],
+            'r152': [3, 8, 36, 3],
+        }
+        block_units = block_units_dict.get(model_name, [3, 4, 6, 3])
+        model = ResNetV2Model(block_units, width_factor, head_size=out_dim)
+
+        if self.num_channels == 3:
+            flatten_dim = 1024 * width_factor
+        if self.include_conv5:
+            flatten_dim *= 2
+
+        return model, flatten_dim
+
+
+class AxialNetHybridViT(HybridVisionTransformer):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5=False,
+        dropout_rate=0.1,
+        positional_encoding_type="learned",
+        backbone='a50m',
+    ):
+        super(AxialNetHybridViT, self).__init__(
+            img_dim=img_dim,
+            out_dim=out_dim,
+            num_channels=num_channels,
+            embedding_dim=embedding_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            hidden_dim=hidden_dim,
+            include_conv5=include_conv5,
+            dropout_rate=dropout_rate,
+            positional_encoding_type=positional_encoding_type,
+            backbone=backbone,
+        )
+
+    def configure_backbone(self):
+        """
+        Current support offered for all BiT models
+        models from https://github.com/csrhddlam/axial-deeplab/blob/master/lib/models/axialnet.py
+
+        expects model name of style 'a{depth}{width}'
+        where depth in [26, 50, 101]
+        where width in [s, m, l]
+        """
+        backbone = self.backbone
+        out_dim = self.out_dim
+
+        model_name = backbone[:3]
+        width = backbone[-1]
+
+        block_units_dict = {
+            'a26': [1, 2, 4, 1],
+            'a50': [3, 4, 6, 3],
+            'a101': [3, 4, 23, 3],
+        }
+        block_units = block_units_dict.get(model_name, [3, 4, 6, 3])
+
+        scale_factor_dict = {'s': 0.5, 'm': 0.75, 'l': 1.0}
+        scale_factor = scale_factor_dict.get(width, 0.75)
+        model = AxialAttentionNet(
+            block_units, s=scale_factor, num_classes=out_dim
+        )
+
+        if self.num_channels == 3:
+            flatten_dim = int(512 * float(scale_factor / 0.5))
+        if self.include_conv5:
+            flatten_dim *= 2
+
+        return model, flatten_dim

third_party/EgoSpeak/MiniROAD/model/transformer_models/OADTR_ViT_origin.py

@@ -0,0 +1,331 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .decoder import Decoder, DecoderLayer
+from .attn import FullAttention, ProbAttention, AttentionLayer
+from .Transformer import TransformerModel
+from ipdb import set_trace
+from .PositionalEncoding import (
+    FixedPositionalEncoding,
+    LearnedPositionalEncoding,
+)
+
+__all__ = ['ViT_B16', 'ViT_B32', 'ViT_L16', 'ViT_L32', 'ViT_H14']
+
+
+class VisionTransformer_v3(nn.Module):
+    def __init__(
+        self,
+        args,
+        img_dim,
+        patch_dim,
+        out_dim,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        dropout_rate=0.0,
+        attn_dropout_rate=0.0,
+        use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned", with_camera=True, with_motion=True, num_channels=3072,
+    ):
+        super(VisionTransformer_v3, self).__init__()
+
+        assert embedding_dim % num_heads == 0
+        assert img_dim % patch_dim == 0
+        self.with_camera = with_camera
+        self.with_motion = with_motion
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.patch_dim = patch_dim
+        # num_channels = img_dim
+        self.num_channels = num_channels
+        self.dropout_rate = dropout_rate
+        self.attn_dropout_rate = attn_dropout_rate
+        self.conv_patch_representation = conv_patch_representation
+
+        # self.num_patches = int((img_dim // patch_dim) ** 2)
+        self.num_patches = int(img_dim // patch_dim)
+        self.seq_length = self.num_patches + 1
+        self.flatten_dim = patch_dim * patch_dim * num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            embedding_dim,
+            num_layers,
+            num_heads,
+            hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(embedding_dim)
+
+        d_model = args.decoder_embedding_dim
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(embedding_dim + d_model, hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(hidden_dim//2, out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(embedding_dim + d_model, out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+
+        # Decoder
+        factor = 1  # 5
+        dropout = args.decoder_attn_dropout_rate
+        # d_model = args.decoder_embedding_dim
+        n_heads = args.decoder_num_heads
+        d_layers = args.decoder_layers
+        d_ff = args.decoder_embedding_dim_out  # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None
+        activation = 'gelu'  # 'gelu'
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(FullAttention(True, factor, attention_dropout=dropout),  # True
+                                   d_model, n_heads),  # ProbAttention  FullAttention
+                    AttentionLayer(FullAttention(False, factor, attention_dropout=dropout),  # False
+                                   d_model, n_heads),
+                    d_model,
+                    d_ff,
+                    dropout=dropout,
+                    activation=activation,
+                )
+                for l in range(d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.decoder_cls_token = nn.Parameter(torch.zeros(1, args.query_num, d_model))
+        if positional_encoding_type == "learned":
+            self.decoder_position_encoding = LearnedPositionalEncoding(
+                args.query_num, self.embedding_dim, args.query_num
+            )
+        elif positional_encoding_type == "fixed":
+            self.decoder_position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position decoding :', positional_encoding_type)
+        self.classifier = nn.Linear(d_model, out_dim)
+        self.after_dropout = nn.Dropout(p=self.dropout_rate)
+        # self.merge_fc = nn.Linear(d_model, 1)
+        # self.merge_sigmoid = nn.Sigmoid()
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.with_camera and self.with_motion:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.with_camera:
+            x = sequence_input_rgb
+        elif self.with_motion:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1)
+        x = self.position_encoding(x)
+        x = self.pe_dropout(x)   # not delete
+
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # [128, 33, 1024]
+        # x = self.after_dropout(x)  # add
+        # decoder
+        decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1)
+        # decoder_cls_token = self.after_dropout(decoder_cls_token)  # add
+        # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token)  # [128, 8, 1024]
+        dec = self.decoder(decoder_cls_token, x)   # [128, 8, 1024]
+        dec = self.after_dropout(dec)  # add
+        # merge_atte = self.merge_sigmoid(self.merge_fc(dec))  # [128, 8, 1]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)  # [128, 1024]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001)
+        dec_for_token = dec.mean(dim=1)
+        # dec_for_token = dec.max(dim=1)[0]
+        dec_cls_out = self.classifier(dec)
+        # set_trace()
+        # x = self.to_cls_token(x[:, 0])
+        x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1)
+        x = self.mlp_head(x)
+        # x = F.log_softmax(x, dim=-1)
+
+        return x, dec_cls_out
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)
+
+
+def ViT_B16(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 16
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=768,
+        num_heads=12,
+        num_layers=12,
+        hidden_dim=3072,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_B32(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 32
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=768,
+        num_heads=12,
+        num_layers=12,
+        hidden_dim=3072,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_L16(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 16
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1024,
+        num_heads=16,
+        num_layers=24,
+        hidden_dim=4096,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_L32(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 32
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1024,
+        num_heads=16,
+        num_layers=24,
+        hidden_dim=4096,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_H14(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 14
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1280,
+        num_heads=16,
+        num_layers=32,
+        hidden_dim=5120,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )

third_party/EgoSpeak/MiniROAD/model/transformer_models/PositionalEncoding.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+
+class FixedPositionalEncoding(nn.Module):
+    def __init__(self, embedding_dim, max_length=5000):
+        super(FixedPositionalEncoding, self).__init__()
+
+        pe = torch.zeros(max_length, embedding_dim)
+        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, embedding_dim, 2).float()
+            * (-torch.log(torch.tensor(10000.0)) / embedding_dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[: x.size(0), :]
+        return x
+
+
+class LearnedPositionalEncoding(nn.Module):
+    def __init__(self, max_position_embeddings, embedding_dim, seq_length):
+        super(LearnedPositionalEncoding, self).__init__()
+        self.pe = nn.Embedding(max_position_embeddings, embedding_dim)
+        self.seq_length = seq_length
+
+        self.register_buffer(
+            "position_ids",
+            torch.arange(max_position_embeddings).expand((1, -1)),
+        )
+
+    def forward(self, x, position_ids=None):
+        if position_ids is None:
+            position_ids = self.position_ids[:, : self.seq_length]
+
+        try:
+            position_embeddings = self.pe(position_ids)
+            return x + position_embeddings
+        except:
+            import pdb; pdb.set_trace()

third_party/EgoSpeak/MiniROAD/model/transformer_models/Transformer.py

@@ -0,0 +1,82 @@
+from torch import nn
+from .Attention import SelfAttention
+
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(x) + x
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(self.norm(x))
+
+
+class PreNormDrop(nn.Module):
+    def __init__(self, dim, dropout_rate, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(p=dropout_rate)
+        self.fn = fn
+
+    def forward(self, x):
+        return self.dropout(self.fn(self.norm(x)))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout_rate):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(p=dropout_rate),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class TransformerModel(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.1,
+    ):
+        super().__init__()
+        layers = []
+        for _ in range(depth):
+            layers.extend(
+                [
+                    Residual(
+                        PreNormDrop(
+                            dim,
+                            dropout_rate,
+                            SelfAttention(
+                                dim, heads=heads, dropout_rate=attn_dropout_rate
+                            ),
+                        )
+                    ),
+                    Residual(
+                        PreNorm(dim, FeedForward(dim, mlp_dim, dropout_rate))
+                    ),
+                ]
+            )
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.net(x)

third_party/EgoSpeak/MiniROAD/model/transformer_models/ViT.py

@@ -0,0 +1,350 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from model.model_builder import META_ARCHITECTURES as registry
+
+from .attn import AttentionLayer, FullAttention, ProbAttention
+from .decoder import Decoder, DecoderLayer
+from .PositionalEncoding import (FixedPositionalEncoding,
+                                 LearnedPositionalEncoding)
+from .Transformer import TransformerModel
+
+
+# @registry.register('TRANSFORMER')
+@registry.register('OADTR')
+class OADTR(nn.Module):
+    def __init__(self, cfg, use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned"
+    ):
+        super(OADTR, self).__init__()
+        
+        self.img_dim = cfg["window_size"] # img_dim
+        self.out_dim = cfg["num_classes"]
+        
+        self.patch_dim = cfg["patch_dim"]
+        self.query_num = cfg["query_num"]
+        
+        self.embedding_dim = cfg["embedding_dim"]
+        self.num_heads = cfg["num_heads"]
+        self.num_layers = cfg["num_layers"]
+        self.hidden_dim = cfg["hidden_dim"]
+        
+        self.dropout_rate = cfg["dropout"]
+        self.attn_dropout_rate = cfg["attn_dropout_rate"]
+        
+        self.conv_patch_representation = conv_patch_representation
+        
+        self.decoder_embedding_dim = cfg["decoder_embedding_dim"]
+        self.decoder_num_heads = cfg["decoder_num_heads"]
+        self.decoder_layers = cfg["decoder_layers"]
+        self.decoder_embedding_dim_out = cfg["decoder_embedding_dim_out"]
+        self.decoder_attn_dropout_rate = cfg["decoder_attn_dropout_rate"]
+        
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.num_channels = 0
+        if self.use_rgb:
+            self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.num_channels += FEATURE_SIZES[cfg['flow_type']]
+        
+        assert self.embedding_dim % self.num_heads == 0
+        assert self.img_dim % self.patch_dim == 0
+        
+        self.num_patches = int(self.img_dim // self.patch_dim)
+        self.seq_length = self.num_patches + 1
+        # self.seq_length = self.num_patches + self.img_dim
+        self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
+        # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            self.embedding_dim,
+            self.num_layers,
+            self.num_heads,
+            self.hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
+
+        d_model = self.decoder_embedding_dim
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(self.embedding_dim , self.hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dim//2, self.out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(self.embedding_dim + d_model, self.out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+        
+        # Decoder
+        factor = 1  # 5
+        dropout = self.decoder_attn_dropout_rate
+        # d_model = self.decoder_embedding_dim
+        n_heads = self.decoder_num_heads
+        d_layers = self.decoder_layers
+        d_ff = self.decoder_embedding_dim_out  # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None
+        activation = 'gelu'  # 'gelu'
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(FullAttention(True, factor, attention_dropout=dropout),  # True
+                                   d_model, n_heads),  # ProbAttention  FullAttention
+                    AttentionLayer(FullAttention(False, factor, attention_dropout=dropout),  # False
+                                   d_model, n_heads),
+                    d_model,
+                    d_ff,
+                    dropout=dropout,
+                    activation=activation,
+                )
+                for l in range(d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.decoder_cls_token = nn.Parameter(torch.zeros(1, self.query_num, d_model))
+        if positional_encoding_type == "learned":
+            self.decoder_position_encoding = LearnedPositionalEncoding(
+                self.query_num, self.embedding_dim, self.query_num
+            )
+        elif positional_encoding_type == "fixed":
+            self.decoder_position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position decoding :', positional_encoding_type)
+        self.classifier = nn.Linear(d_model, self.out_dim)
+        self.after_dropout = nn.Dropout(p=self.dropout_rate)
+        # self.merge_fc = nn.Linear(d_model, 1)
+        # self.merge_sigmoid = nn.Sigmoid()
+
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.use_rgb and self.use_flow:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.use_rgb:
+            x = sequence_input_rgb
+        elif self.use_flow:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x) # B, 128, 1204
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
+        try:
+            x = self.position_encoding(x) # B, seq+1, 1024
+        except:
+            import pdb; pdb.set_trace()
+        x = self.pe_dropout(x)   # not delete
+        
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # B, seq+1, 1024
+
+        # decoder
+        decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1) # [B, 8, 1024]
+        # decoder_cls_token = self.after_dropout(decoder_cls_token)  # add
+        # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token)  # [128, 8, 1024]
+        dec = self.decoder(decoder_cls_token, x)   # [128, 8, 1024]
+        dec = self.after_dropout(dec)  # add
+        
+        # merge_atte = self.merge_sigmoid(self.merge_fc(dec))  # [128, 8, 1]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)  # [128, 1024]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001)
+        # *** RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1024x2)
+        dec_for_token = dec.mean(dim=1) # [64, 1024]
+        
+        # dec_for_token = dec.max(dim=1)[0]
+        dec_cls_out = self.classifier(dec) # [B, 8, NUM_CLASS?]
+        # set_trace()
+        # x = self.to_cls_token(x[:, 0])
+        
+        x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1) # [B, 2048]
+        x = self.mlp_head(x) # [B, NUM_CLASSES]
+        
+        # original OADTR
+        # return x, dec_cls_out
+        
+        # x = F.log_softmax(x, dim=-1)
+        out_dict = {}
+        out_dict['logits'] = x.unsqueeze(1) # [B, 1, NUM_CLASSES] ex) [64, 1, 2]
+        out_dict['dec_logits'] = dec_cls_out # [B, 8, NUM_CLASSES] ex) [64, 8, 2]
+        
+        # out_dict['logits'] = x #x.unsqueeze(1)
+        return out_dict
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)
+
+@registry.register('Transformer')
+class ViTEnc(nn.Module):
+    def __init__(self, cfg, use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned"
+    ):
+        super(ViTEnc, self).__init__()
+        self.img_dim = cfg["window_size"]
+        self.out_dim = cfg["num_classes"]
+        self.embedding_dim = cfg["embedding_dim"]
+        self.patch_dim = cfg["patch_dim"]
+        self.num_heads = cfg["num_heads"]
+        self.num_layers = cfg["num_layers"]
+        self.hidden_dim = cfg["hidden_dim"]
+        self.dropout_rate = cfg["dropout"]
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.num_channels= 0
+        if self.use_rgb:
+            self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.num_channels += FEATURE_SIZES[cfg['flow_type']]
+        self.attn_dropout_rate = cfg["attn_dropout_rate"]
+        assert self.embedding_dim % self.num_heads == 0
+        assert self.img_dim % self.patch_dim == 0
+        self.conv_patch_representation = conv_patch_representation
+        self.num_patches = int(self.img_dim // self.patch_dim)
+        self.seq_length = self.num_patches + 1
+        # self.seq_length = self.num_patches + self.img_dim
+        self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
+        # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            self.embedding_dim,
+            self.num_layers,
+            self.num_heads,
+            self.hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
+
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(self.embedding_dim , self.hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dim//2, self.out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(self.embedding_dim, self.out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.use_rgb and self.use_flow:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.use_rgb:
+            x = sequence_input_rgb
+        elif self.use_flow:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
+        
+        x = self.position_encoding(x) # B, seq+1, 1024
+        x = self.pe_dropout(x)   # not delete
+
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # B, seq+1, 1024
+
+        x = self.to_cls_token(x[:, 0]) # B, 1024
+        # x = self.to_cls_token(x[:,0:self.img_dim]) # B, 1024
+        x = self.mlp_head(x)
+        # x = F.log_softmax(x, dim=-1)
+        out_dict = {}
+        out_dict['logits'] = x.unsqueeze(1)
+        # out_dict['logits'] = x #x.unsqueeze(1)
+        return out_dict
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)

third_party/EgoSpeak/MiniROAD/model/transformer_models/__init__.py

@@ -0,0 +1,4 @@
+from .ViT import ViTEnc
+from .HybridViT import ResNetHybridViT, AxialNetHybridViT
+
+__all__ = ['ResNetHybridViT', 'AxialNetHybridViT', 'VisionTransformer_v3']

third_party/EgoSpeak/MiniROAD/model/transformer_models/attn.py

@@ -0,0 +1,170 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+
+from math import sqrt
+
+
+class TriangularCausalMask():
+    def __init__(self, B, L, device="cpu"):
+        mask_shape = [B, 1, L, L]
+        with torch.no_grad():
+            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class ProbMask():
+    def __init__(self, B, H, L, index, scores, device="cpu"):
+        _mask = torch.ones(L, scores.shape[-1], dytpe=torch.bool).to(device).triu(1)
+        _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])
+        indicator = _mask_ex[torch.arange(B)[:, None, None],
+                    torch.arange(H)[None, :, None],
+                    index, :].to(device)
+        self._mask = indicator.view(scores.shape).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class FullAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1):
+        super(FullAttention, self).__init__()
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        scale = self.scale or 1. / sqrt(E)
+
+        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+        if self.mask_flag:
+            if attn_mask is None:
+                attn_mask = TriangularCausalMask(B, L, device=queries.device)
+
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        A = self.dropout(torch.softmax(scale * scores, dim=-1))
+        V = torch.einsum("bhls,bshd->blhd", A, values)
+
+        return V.contiguous()
+
+
+class ProbAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1):
+        super(ProbAttention, self).__init__()
+        self.factor = factor
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def _prob_QK(self, Q, K, sample_k, n_top):
+        # Q [B, H, L, D]
+        B, H, L, E = K.shape
+        _, _, S, _ = Q.shape
+
+        # calculate the sampled Q_K
+        K_expand = K.unsqueeze(-3).expand(B, H, S, L, E)
+        indx_sample = torch.randint(L, (S, sample_k))
+        K_sample = K_expand[:, :, torch.arange(S).unsqueeze(1), indx_sample, :]
+        Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
+
+        # find the Top_k query with sparisty measurement
+        M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L)
+        M_top = M.topk(n_top, sorted=False)[1]
+
+        # use the reduced Q to calculate Q_K
+        Q_reduce = Q[torch.arange(B)[:, None, None],
+                   torch.arange(H)[None, :, None],
+                   M_top, :]
+        Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1))
+
+        return Q_K, M_top
+
+    def _get_initial_context(self, V, L_Q):
+        B, H, L_V, D = V.shape
+        if not self.mask_flag:
+            V_sum = V.sum(dim=-2)
+            contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
+        else:  # use mask
+            assert (L_Q == L_V)  # requires that L_Q == L_V, i.e. for self-attention only
+            contex = V.cumsum(dim=-1)
+        return contex
+
+    def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
+        B, H, L_V, D = V.shape
+
+        if self.mask_flag:
+            attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        attn = torch.softmax(scores, dim=-1)  # nn.Softmax(dim=-1)(scores)
+
+        context_in[torch.arange(B)[:, None, None],
+        torch.arange(H)[None, :, None],
+        index, :] = torch.matmul(attn, V)
+        return context_in
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, D = queries.shape
+        _, S, _, _ = keys.shape
+
+        queries = queries.view(B, H, L, -1)
+        keys = keys.view(B, H, S, -1)
+        values = values.view(B, H, S, -1)
+
+        U = self.factor * np.ceil(np.log(S)).astype('int').item()
+        u = self.factor * np.ceil(np.log(L)).astype('int').item()
+
+        scores_top, index = self._prob_QK(queries, keys, u, U)
+        # add scale factor
+        scale = self.scale or 1. / sqrt(D)
+        if scale is not None:
+            scores_top = scores_top * scale
+        # get the context
+        context = self._get_initial_context(values, L)
+        # update the context with selected top_k queries
+        context = self._update_context(context, values, scores_top, index, L, attn_mask)
+
+        return context.contiguous()
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, attention, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AttentionLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_attention = attention
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out = self.inner_attention(
+            queries,
+            keys,
+            values,
+            attn_mask
+        ).view(B, L, -1)
+
+        return self.out_projection(out)

third_party/EgoSpeak/MiniROAD/model/transformer_models/decoder.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
+                 dropout=0.1, activation="relu"):
+        super(DecoderLayer, self).__init__()
+        d_ff = d_ff or 4*d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask
+        ))
+        x = self.norm1(x)
+
+        x = x + self.dropout(self.cross_attention(
+            x, cross, cross,
+            attn_mask=cross_mask
+        ))
+
+        y = x = self.norm2(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1,1))))
+        y = self.dropout(self.conv2(y).transpose(-1,1))
+
+        return self.norm3(x+y)
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, norm_layer=None):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        for layer in self.layers:
+            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x

third_party/EgoSpeak/MiniROAD/model/weights_init.py

@@ -0,0 +1,35 @@
+# FROM LSTR
+import math
+import torch.nn as nn
+
+def weights_init(m):
+    if isinstance(m, nn.Linear):
+        nn.init.kaiming_uniform_(m.weight.data, a=math.sqrt(5))
+    elif isinstance(m, nn.Conv1d):
+        nn.init.normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv2d):
+        nn.init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose1d):
+        nn.init.normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose2d):
+        nn.init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.BatchNorm1d):
+        nn.init.normal_(m.weight.data, mean=1, std=0.02)
+        nn.init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm2d):
+        nn.init.normal_(m.weight.data, mean=1, std=0.02)
+        nn.init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.GRUCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                nn.init.orthogonal_(param.data)
+            else:
+                nn.init.normal_(param.data)

third_party/EgoSpeak/MiniROAD/run_mamba_inference.sh

@@ -0,0 +1,26 @@
+PATH_TO_CONFIG_FILE=configs/EasyCom/mamba_aa_easycom.yaml
+PATH_TO_CHECKPOINT_FILE=checkpoints/mamba_easycom_AV.pth
+shift  # Remove the first argument so $@ contains all other arguments
+
+echo "Running MiniROAD inference"
+echo "Config file: $PATH_TO_CONFIG_FILE"
+echo "Additional options: $@"
+
+# get the basename of the config file, without the extension
+CONFIG_FILENAME=$(basename -- "$PATH_TO_CONFIG_FILE")
+CONFIG_FILENAME="${CONFIG_FILENAME%.*}"
+
+# --no_flow (for no audio) --no_rgb --wandb
+# Pass all remaining arguments directly to the Python script
+python -u main.py --config "$PATH_TO_CONFIG_FILE" \
+      --save-as-cfg \
+      --eval \
+      --ckpt "$PATH_TO_CHECKPOINT_FILE" \
+      $@
+
+echo "Finished MiniROAD inference"
+echo "python -u main.py --config '$PATH_TO_CONFIG_FILE' \
+      --save-as-cfg \
+      --ckpt '$PATH_TO_CHECKPOINT_FILE' \
+      --eval \
+      $@"

third_party/EgoSpeak/MiniROAD/run_miniroad_inference.sh

@@ -0,0 +1,26 @@
+PATH_TO_CONFIG_FILE=configs/EasyCom/miniroad_aa_easycom.yaml
+PATH_TO_CHECKPOINT_FILE=checkpoints/miniroad_easycom_AV.pth
+shift  # Remove the first argument so $@ contains all other arguments
+
+echo "Running MiniROAD inference"
+echo "Config file: $PATH_TO_CONFIG_FILE"
+echo "Additional options: $@"
+
+# get the basename of the config file, without the extension
+CONFIG_FILENAME=$(basename -- "$PATH_TO_CONFIG_FILE")
+CONFIG_FILENAME="${CONFIG_FILENAME%.*}"
+
+# --no_flow (for no audio) --no_rgb --wandb
+# Pass all remaining arguments directly to the Python script
+python -u main.py --config "$PATH_TO_CONFIG_FILE" \
+      --save-as-cfg \
+      --eval \
+      --ckpt "$PATH_TO_CHECKPOINT_FILE" \
+      $@
+
+echo "Finished MiniROAD inference"
+echo "python -u main.py --config '$PATH_TO_CONFIG_FILE' \
+      --save-as-cfg \
+      --ckpt '$PATH_TO_CHECKPOINT_FILE' \
+      --eval \
+      $@"

third_party/EgoSpeak/MiniROAD/run_miniroad_train.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+PATH_TO_CONFIG_FILE=configs/EasyCom/miniroad_aa_easycom.yaml
+
+echo "Running MiniROAD training"
+echo "Config file: $PATH_TO_CONFIG_FILE"
+echo "Additional options: $@"
+
+# Additional options: --no_flow (for no audio) --no_rgb --wandb
+# Pass all remaining arguments directly to the Python script
+python -u main.py --config "$PATH_TO_CONFIG_FILE" --save-as-cfg $@
+
+echo "Finished MiniROAD training"
+echo "Run python -u main.py --config $PATH_TO_CONFIG_FILE --save-as-cfg $@"

third_party/EgoSpeak/MiniROAD/trainer/__init__.py

@@ -0,0 +1,4 @@
+from .train_builder import build_trainer
+from .train import train_one_epoch
+from .eval_builder import build_eval
+from .eval import Evaluate

third_party/EgoSpeak/MiniROAD/trainer/eval.py

@@ -0,0 +1,288 @@
+import json
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import wandb
+from tqdm import tqdm
+from trainer.eval_builder import EVAL
+from utils import *
+from utils import (event_accuracy, perframe_average_precision,
+                   thumos_postprocessing)
+
+
+@EVAL.register("OAD")
+class Evaluate(nn.Module):
+    
+    def __init__(self, cfg):
+        super(Evaluate, self).__init__()
+        self.data_processing = thumos_postprocessing if 'THUMOS' in cfg['data_name'] else None
+        self.metric = cfg['metric']
+        self.eval_method = perframe_average_precision
+        self.event_eval_method = event_accuracy
+        self.all_class_names = json.load(open(cfg['video_list_path']))[cfg["data_name"].split('_')[0]]['class_index']
+        self.cfg = cfg
+        self.ignore_index = cfg['ignore_index'] if 'ignore_index' in cfg else -100
+    
+    def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path):
+        device = "cuda:0"
+        model.eval()
+        eval_preds = {}
+        gt_targets_dict = {}
+        
+        with torch.no_grad():
+            pred_scores, gt_targets = [], []
+            start = time.time()
+            eval_loss = 0.0
+            for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(dataloader, desc='Evaluation:', leave=False)):
+                rgb_input, flow_input, target = rgb_input.to(device), flow_input.to(device), target.to(device)
+                out_dict = model(rgb_input, flow_input)
+                
+                if criterion != None:
+                    loss = criterion(out_dict, target)
+                    eval_loss += loss.item()
+                    
+                pred_logit = out_dict['logits']
+                prob_val = pred_logit.squeeze().cpu().numpy()
+                target_batch = target.squeeze().cpu().numpy()
+                pred_scores += list(prob_val) 
+                gt_targets += list(target_batch)
+                
+                eval_preds[vid_name] = prob_val
+    
+            end = time.time()
+            num_frames = len(gt_targets)
+            
+            # save prediction scores
+            # if result_path != None:
+            time_taken = end - start
+            logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)')
+                
+            result_AP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, 'AP')
+            result_cAP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, 'cAP')
+            logger.info(f'per_class_ap: {result_AP["per_class_AP"]}')
+            logger.info(f'all_class_names: {self.all_class_names}')
+            logger.info(f'len(result_AP["per_class_AP"]: {len(result_AP["per_class_AP"])}')
+            
+            per_class_AP = result_AP["per_class_AP"]
+            if use_wandb == True:
+                wandb.log({"Validation Loss": eval_loss / len(dataloader)})
+                wandb.log({"Validation per_class_AP": per_class_AP})            
+                
+            if writer != None:
+                writer.add_scalar("Validation Loss", eval_loss / len(dataloader), epoch)
+            
+            sum_class_AP = 0.0
+            
+            for class_name in self.all_class_names:
+                # TVSeries
+                if class_name == 'background':
+                    continue
+                # THUMOS
+                if class_name == 'Background' or class_name == 'Ambiguous':
+                    continue
+                
+                # logger.info(f'{class_name}')
+                class_AP = result_AP["per_class_AP"][class_name]
+                sum_class_AP += class_AP
+                
+                class_cAP = result_cAP["per_class_AP"][class_name]
+                logger.info(f'AP @ {class_name}: {class_AP:.4f}') 
+                logger.info(f'cAP @ {class_name}: {class_cAP:.4f}')
+
+            logger.info(f'sum_AP: {sum_class_AP}, mean_AP: {sum_class_AP / (len(self.all_class_names) - 1):.4f}')
+            
+            logger.info(f'pred_scores: {pred_scores[:2]}')
+            logger.info(f'len(pred_scores): {len(pred_scores)}')
+            logger.info(f'len(pred_scores[0]): {len(pred_scores[0])}')
+            
+        return result_AP['mean_AP'], result_cAP['mean_AP'], eval_preds
+    
+    def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None, result_path = None):
+        return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path)
+
+@EVAL.register("OADTR")
+class EvaluateOADTR(nn.Module):
+    
+    def __init__(self, cfg):
+        super(EvaluateOADTR, self).__init__()
+        self.data_processing = thumos_postprocessing if 'THUMOS' in cfg['data_name'] else None
+        self.metric = cfg['metric']
+        self.eval_method = perframe_average_precision
+        self.all_class_names = json.load(open(cfg['video_list_path']))[cfg["data_name"].split('_')[0]]['class_index']
+        self.cfg = cfg
+    
+    def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer):
+        device = "cuda:0"
+        model.eval()
+        eval_preds = {}
+        gt_targets_dict = {}
+        
+        with torch.no_grad():
+            pred_scores, gt_targets = [], []
+            start = time.time()
+            eval_loss = 0.0
+            for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(dataloader, desc='Evaluation:', leave=False)):
+                rgb_input, flow_input, target = rgb_input.to(device), flow_input.to(device), target.to(device)
+                out_dict = model(rgb_input, flow_input)
+                
+                if criterion != None:
+                    loss = criterion(out_dict, target)
+                    eval_loss += loss.item()
+                    
+                pred_logit = out_dict['logits']
+                prob_val = pred_logit.squeeze().cpu().numpy()
+                target_batch = target.squeeze().cpu().numpy()
+                pred_scores += list(prob_val) 
+                gt_targets += list(target_batch)
+                
+                # eval_preds[vid_name] = pred_scores
+                # gt_targets_dict[vid_name] = pred_scores
+            
+            if use_wandb == True:
+                wandb.log({"Validation Loss": eval_loss / len(dataloader)})
+                
+            if writer != None:
+                writer.add_scalar("Validation Loss", eval_loss / len(dataloader), epoch)
+            
+            end = time.time()
+            num_frames = len(gt_targets)
+            result_AP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.data_processing, 'AP')
+            result_cAP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.data_processing, 'cAP')
+            time_taken = end - start
+            logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)')
+            # logger.info(f'per_class_ap: {result["per_class_AP"]}')
+            logger.info(f'all_class_names: {self.all_class_names}')
+            logger.info(f'len(result_AP["per_class_AP"]: {len(result_AP["per_class_AP"])}')
+            
+            sum_class_AP = 0.0
+            
+            for class_name in self.all_class_names:
+                # TVSeries
+                if class_name == 'background':
+                    continue
+                # THUMOS
+                if class_name == 'Background' or class_name == 'Ambiguous':
+                    continue
+                
+                # logger.info(f'{class_name}')
+                class_AP = result_AP["per_class_AP"][class_name]
+                sum_class_AP += class_AP
+                
+                class_cAP = result_cAP["per_class_AP"][class_name]
+                logger.info(f'AP @ {class_name}: {class_AP:.4f}') 
+                logger.info(f'cAP @ {class_name}: {class_cAP:.4f}')
+
+            logger.info(f'sum_AP: {sum_class_AP}, mean_AP: {sum_class_AP / (len(self.all_class_names) - 1):.4f}')
+            
+            logger.info(f'pred_scores: {pred_scores[:2]}')
+            logger.info(f'len(pred_scores): {len(pred_scores)}')
+            logger.info(f'len(pred_scores[0]): {len(pred_scores[0])}')
+            
+        return result_AP['mean_AP'], result_cAP['mean_AP'], eval_preds
+    
+    def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None):
+        return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer)
+
+
+@EVAL.register("ANTICIPATION")
+class ANT_Evaluate(nn.Module):
+    
+    def __init__(self, cfg):
+        super(ANT_Evaluate, self).__init__()
+        data_name = cfg["data_name"].split('_')[0]
+        self.data_processing = thumos_postprocessing if data_name == 'THUMOS' else None
+        self.metric = cfg['metric']
+        self.eval_method = perframe_average_precision
+        self.event_eval_method = event_accuracy
+        self.all_class_names = json.load(open(cfg['video_list_path']))[data_name]['class_index']
+        self.cfg = cfg
+        
+        if 'ignore_index' in cfg:
+            self.ignore_index = cfg['ignore_index']
+        else:
+            raise NotImplementedError("Please set ignore_index in config")
+            
+    def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path):
+        device = "cuda:0"
+        model.eval()
+        ant_eval_pred_scores_dict = {}
+        ant_gt_targets_dict = {}
+        
+        with torch.no_grad():
+            pred_scores, gt_targets, ant_pred_scores, ant_gt_targets = [], [], [], []
+            start = time.time()
+            anticipation_mAPs = []
+            for vid_name, rgb_input, flow_input, target, ant_target in tqdm(dataloader, desc='Evaluation:', leave=False):
+                rgb_input, flow_input, target, ant_target = rgb_input.to(device), flow_input.to(device), target.to(device), ant_target.to(device)
+                
+                out_dict = model(rgb_input, flow_input)                
+                pred_logit = out_dict['logits']
+                ant_pred_logit = out_dict['anticipation_logits']
+                
+                prob_val = pred_logit.squeeze().cpu().numpy() # [8390, 3]
+                target_batch = target.squeeze().cpu().numpy()
+                ant_prob_val = ant_pred_logit.squeeze().cpu().numpy() # [8390, ANTICIPATION_LENGTH, 3]
+                ant_target_batch = ant_target.squeeze().cpu().numpy()
+                
+                pred_scores += list(prob_val)
+                gt_targets += list(target_batch)
+                ant_pred_scores += list(ant_prob_val)
+                ant_gt_targets += list(ant_target_batch)
+                if isinstance(vid_name, tuple):
+                    vid_name = vid_name[0]
+                    
+                ant_eval_pred_scores_dict[vid_name] = ant_prob_val
+                ant_gt_targets_dict[vid_name] = ant_target_batch
+
+            end = time.time()
+            num_frames = len(gt_targets)
+            
+            result = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, self.metric)
+            per_class_AP = result['per_class_AP']
+            
+            ant_pred_scores = np.array(ant_pred_scores)
+            ant_gt_targets = np.array(ant_gt_targets)
+            logger.info(f'OAD mAP: {result["mean_AP"]*100:.2f}')
+            # Prepare the formatted string
+            formatted_output = "per_class_AP:\n" + "\n".join(f"  {class_name:<25} : {ap_value:.4f}" for class_name, ap_value in per_class_AP.items())
+
+            # Log the formatted string
+            logger.info(formatted_output)
+            
+            if use_wandb == True:
+                wandb.log({"OAD mAP": result["mean_AP"]})
+            
+            anticipation_mAPs_per_class = {key: [] for key in per_class_AP.keys()}
+            
+            for step in range(ant_gt_targets.shape[1]):                
+                result[f'anticipation_{step+1}'] = self.eval_method(ant_pred_scores[:,step,:], ant_gt_targets[:,step,:], self.all_class_names, self.ignore_index, self.data_processing, self.metric)
+                anticipation_mAPs.append(result[f'anticipation_{step+1}']['mean_AP'])
+                formatted_output = "per_class_AP:\n" + "\n".join(f"  {class_name:<25} : {ap_value:.4f}" for class_name, ap_value in result[f'anticipation_{step+1}']['per_class_AP'].items())
+                logger.info(f"Anticipation at step {step+1}: {result[f'anticipation_{step+1}']['mean_AP']*100:.2f}")
+                logger.info(formatted_output)
+                
+                for class_name in self.all_class_names:
+                    anticipation_mAPs_per_class[class_name].append(result[f'anticipation_{step+1}']['per_class_AP'][class_name])
+                
+                if use_wandb == True:
+                    wandb.log({f"Anticipation at step {step+1}": result[f'anticipation_{step+1}']['mean_AP']})
+            
+            logger.info(f'Mean Anticipation mAP: {np.mean(anticipation_mAPs)*100:.2f}')
+            formatted_output = "Mean Anticipation per_class_AP:\n" + "\n".join(f"  {class_name:<25} : {np.mean(anticipation_mAPs_per_class[class_name])*100:.4f}" for class_name in self.all_class_names)
+            # logger.info(f'Mean Anticipation mAP per class: {np.mean(list(anticipation_mAPs_per_class.values()), axis=1)}')
+            logger.info(formatted_output)
+
+            if use_wandb == True:
+                wandb.log({"Mean Anticipation mAP": np.mean(anticipation_mAPs)})
+            
+            time_taken = end - start
+            logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)')
+        
+        return np.mean(anticipation_mAPs), np.mean(anticipation_mAPs), ant_eval_pred_scores_dict
+    
+    # def forward(self, model, dataloader, logger):
+    #     return self.eval(model, dataloader, logger)
+    def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None, result_path=None):
+        return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path)

third_party/EgoSpeak/MiniROAD/trainer/eval_builder.py

@@ -0,0 +1,11 @@
+__all__ = [
+    'build_eval'
+]
+
+from utils import Registry
+
+EVAL = Registry()
+
+def build_eval(cfg):
+    eval = EVAL[cfg["task"]](cfg)
+    return eval

third_party/EgoSpeak/MiniROAD/trainer/train.py

@@ -0,0 +1,96 @@
+import torch
+import wandb
+from tqdm import tqdm
+from trainer.train_builder import TRAINER
+
+
+@TRAINER.register("OAD")
+def train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, target = rgb_input.cuda(), flow_input.cuda(), target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():    
+                out_dict = model(rgb_input, flow_input) 
+                loss = criterion(out_dict, target)
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            # scaler is None as argparse by default
+        else:
+            out_dict = model(rgb_input, flow_input)
+            loss = criterion(out_dict, target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+        
+        epoch_loss += loss.item()
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+    return epoch_loss
+
+@TRAINER.register("OADTR")
+def train_one_epoch_oadtr(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, enc_target, distance_target, class_h_target = rgb_input.cuda(), flow_input.cuda(), enc_target.cuda(), distance_target.cuda(), class_h_target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():    
+                out_dict = model(rgb_input, flow_input) 
+                loss = criterion(out_dict, enc_target, dec_target)   
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            out_dict = model(rgb_input, flow_input) 
+            loss = criterion(out_dict, enc_target, dec_target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+
+        epoch_loss += loss.item()
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+            
+    return epoch_loss
+
+
+@TRAINER.register("ANTICIPATION")
+def ant_train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, target, ant_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), target.cuda(), ant_target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():
+                out_dict = model(rgb_input, flow_input)
+                loss = criterion(out_dict, target, ant_target)
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            out_dict = model(rgb_input, flow_input) 
+            loss = criterion(out_dict, target, ant_target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+        epoch_loss += loss.item()
+        
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+            
+    return epoch_loss

third_party/EgoSpeak/MiniROAD/trainer/train_builder.py

@@ -0,0 +1,11 @@
+__all__ = [
+    'build_trainer'
+]
+
+from utils import Registry
+
+TRAINER = Registry()
+
+def build_trainer(cfg):
+    trainer = TRAINER[cfg["task"]]
+    return trainer

third_party/EgoSpeak/MiniROAD/utils/__init__.py

@@ -0,0 +1,7 @@
+from .metrics import perframe_average_precision, perstage_average_precision, event_accuracy
+from .postprocessing import thumos_postprocessing
+from .util import *
+from .group_transforms import *
+from .lr_scheduler import build_lr_scheduler
+from .logger import get_logger
+from .registry import Registry

third_party/EgoSpeak/MiniROAD/utils/group_transforms.py

@@ -0,0 +1,497 @@
+"""
+This file is inherited from tsn-pytorch
+"""
+
+import torchvision
+import random
+from PIL import Image, ImageOps
+import numpy as np
+import numbers
+import math
+import torch
+from torchvision.transforms import functional as F
+
+# lstr
+class Compose(object):
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, data):
+        for t in self.transforms:
+            data = t(data)
+        return data
+
+#LSTR
+class GroupCrop(object):
+    """perfrom group crop'"""
+
+    def __init__(self, crop_size, crop_pos='random'):
+        super(GroupCrop, self).__init__()
+        self.crop_size = crop_size
+        self.crop_pos = crop_pos
+        assert crop_pos in ['random', 'top_left', 'center', 'bottom_right'
+                           ], f'crop_pos: {self.crop_pos} not implemented'
+
+    def __call__(self, img_group):
+        """
+        Args:
+            img_group (List(PIL.Image) or List(torch.Tensor)): a list of image to crop.
+        Returns: The same type as input. cropped img_group.
+            
+        """
+        crop_w, crop_h = self.crop_size
+        if isinstance(img_group, torch.Tensor):
+            assert img_group.dim(
+            ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}'
+            img_h, img_w = img_group.shape[-2:]
+            offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w,
+                                                        crop_h)
+            return F.crop(img_group, offset_h, offset_w, crop_h, crop_w)
+        elif isinstance(img_group[0], Image):
+            img_w, img_h = img_group[0].size
+            offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w,
+                                                        crop_h)
+            return [
+                F.crop(img, offset_h, offset_w, crop_h, crop_w)
+                for img in img_group
+            ]
+        else:
+            raise NotImplementedError(
+                f'img_group of type: {type(img_group)} is not supported')
+
+    def get_crop_offset(self, img_w, img_h, crop_w, crop_h):
+        """get the offset to crop.
+        Returns: (offset_h, offset_w).
+        """
+        if self.crop_pos == 'random':
+            offset_h = int(torch.randint(img_h - crop_h, []))
+            offset_w = int(torch.randint(img_w - crop_w, []))
+        elif self.crop_pos == 'top_left':
+            offset_h = offset_w = 0
+        elif self.crop_pos == 'center':
+            offset_h = int((img_h - crop_h) // 2)
+            offset_w = int((img_w - crop_w) // 2)
+        elif self.crop_pos == 'bottom_right':
+            offset_h = img_h - crop_h
+            offset_w = img_w - crop_w
+        else:
+            raise NotImplementedError(
+                f'crop_pos: {self.crop_pos} not implemented')
+        return (offset_h, offset_w)
+
+class GroupRandomCrop(object):
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, img_group):
+
+        w, h = img_group[0].size
+        th, tw = self.size
+
+        out_images = list()
+
+        x1 = random.randint(0, w - tw)
+        y1 = random.randint(0, h - th)
+
+        for img in img_group:
+            assert(img.size[0] == w and img.size[1] == h)
+            if w == tw and h == th:
+                out_images.append(img)
+            else:
+                out_images.append(img.crop((x1, y1, x1 + tw, y1 + th)))
+
+        return out_images
+
+
+class GroupCenterCrop(object):
+    def __init__(self, size):
+        self.worker = torchvision.transforms.CenterCrop(size)
+
+    def __call__(self, img_group):
+        return [self.worker(img) for img in img_group]
+
+
+class GroupRandomHorizontalFlip(object):
+    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
+    """
+    def __init__(self, is_flow=False):
+        self.is_flow = is_flow
+
+    def __call__(self, img_group, is_flow=False):
+        v = random.random()
+        if v < 0.5:
+            ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group]
+            if self.is_flow:
+                for i in range(0, len(ret), 2):
+                    ret[i] = ImageOps.invert(ret[i])  # invert flow pixel values when flipping
+            return ret
+        else:
+            return img_group
+
+
+# class GroupNormalize(object):
+#     def __init__(self, mean, std):
+#         self.mean = mean
+#         self.std = std
+
+#     def __call__(self, tensor):
+#         rep_mean = self.mean * (tensor.size()[0]//len(self.mean))
+#         rep_std = self.std * (tensor.size()[0]//len(self.std))
+
+#         # TODO: make efficient
+#         for t, m, s in zip(tensor, rep_mean, rep_std):
+#             t.sub_(m).div_(s)
+
+#         return tensor
+
+# lstr
+class GroupNormalize(object):
+
+    def __init__(self, mean, std):
+
+        if len(mean) == 1:    # flow
+            assert len(
+                std
+            ) == 1, f'flow input must have std of length 1. got {len(std)}'
+            self.mean = mean[0]
+            self.std = std[0]
+        else:
+            mean = np.array(mean, dtype=np.float32)
+            self.mean = mean[np.newaxis, :, np.newaxis,
+                                np.newaxis]    # shape: (C,) ->(T,C,H,W)
+            std = np.array(std, dtype=np.float32)
+            self.std = std[np.newaxis, :, np.newaxis,
+                            np.newaxis]    # shape: (C,) ->(T,C,H,W)
+
+    def __call__(self, tensor):
+        """
+        Args:
+            tensor (torch.Tensor) : The stacked images. shape: (T, C, H, W).
+        Returns: (torch.Tensor).  sub mean and div std along the channel dim.
+            
+        """
+        return (tensor - self.mean) / self.std
+
+class GroupScale(object):
+    """ Rescales the input PIL.Image to the given 'size'.
+    'size' will be the size of the smaller edge.
+    For example, if height > width, then image will be
+    rescaled to (size * height / width, size)
+    size: size of the smaller edge
+    interpolation: Default: PIL.Image.BILINEAR
+    """
+
+    def __init__(self, size, interpolation=Image.BILINEAR):
+        self.worker = torchvision.transforms.Resize(size, interpolation)
+
+    def __call__(self, img_group):
+        return [self.worker(img) for img in img_group]
+
+# lstr
+class GroupResize(object):
+    """perfrom group random resize"""
+
+    def __init__(self, scale_range, keep_aspect_ratio=True):
+        """
+        Args:
+            scale_range (tuple/list): The shorter side of the image will be resized to
+                                        1) the range if `scale_range` is a tuple
+                                        2) the specific size if `scale_range` contains only 1 element.
+        """
+        super(GroupResize, self).__init__()
+        self.scale_range = scale_range
+        self.keep_aspect_ratio = keep_aspect_ratio
+
+    def __call__(self, img_group):
+        """
+        Args:
+            img_group (List(PIL.Image) or Tensor([T,C,H,W])): a list of image to resize.
+        Returns: The same type as input. resized img_group.
+            
+        """
+        # resize
+        size = self.scale_range[0] if len(
+            self.scale_range) == 1 else torch.randint(self.scale_range[0],
+                                                        self.scale_range[1], [])
+        size = int(size)
+        size = size if self.keep_aspect_ratio else (size, size)
+
+        if isinstance(img_group, torch.Tensor):
+            assert img_group.dim(
+            ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}'
+            return F.resize(img_group, size)
+        elif isinstance(img_group[0], Image):
+            return [F.resize(img, size) for img in img_group]
+        else:
+            raise NotImplementedError(
+                f'img_group of type: {type(img_group)} is not supported')
+
+class GroupOverSample(object):
+    def __init__(self, crop_size, scale_size=None):
+        self.crop_size = crop_size if not isinstance(crop_size, int) else (crop_size, crop_size)
+
+        if scale_size is not None:
+            self.scale_worker = GroupScale(scale_size)
+        else:
+            self.scale_worker = None
+
+    def __call__(self, img_group):
+
+        if self.scale_worker is not None:
+            img_group = self.scale_worker(img_group)
+        image_w, image_h = img_group[0].size
+        crop_w, crop_h = self.crop_size
+
+        offsets = GroupMultiScaleCrop.fill_fix_offset(False, image_w, image_h, crop_w, crop_h)
+        oversample_group = list()
+        for o_w, o_h in offsets:
+            normal_group = list()
+            flip_group = list()
+            for i, img in enumerate(img_group):
+                crop = img.crop((o_w, o_h, o_w + crop_w, o_h + crop_h))
+                normal_group.append(crop)
+                flip_crop = crop.copy().transpose(Image.FLIP_LEFT_RIGHT)
+
+                if img.mode == 'L' and i % 2 == 0:
+                    flip_group.append(ImageOps.invert(flip_crop))
+                else:
+                    flip_group.append(flip_crop)
+
+            oversample_group.extend(normal_group)
+            oversample_group.extend(flip_group)
+        return oversample_group
+
+
+class GroupMultiScaleCrop(object):
+
+    def __init__(self, input_size, scales=None, max_distort=1, fix_crop=True, more_fix_crop=True):
+        self.scales = scales if scales is not None else [1, 875, .75, .66]
+        self.max_distort = max_distort
+        self.fix_crop = fix_crop
+        self.more_fix_crop = more_fix_crop
+        self.input_size = input_size if not isinstance(input_size, int) else [input_size, input_size]
+        self.interpolation = Image.BILINEAR
+
+    def __call__(self, img_group):
+
+        im_size = img_group[0].size
+
+        crop_w, crop_h, offset_w, offset_h = self._sample_crop_size(im_size)
+        crop_img_group = [img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h)) for img in img_group]
+        ret_img_group = [img.resize((self.input_size[0], self.input_size[1]), self.interpolation)
+                         for img in crop_img_group]
+        return ret_img_group
+
+    def _sample_crop_size(self, im_size):
+        image_w, image_h = im_size[0], im_size[1]
+
+        # find a crop size
+        base_size = min(image_w, image_h)
+        crop_sizes = [int(base_size * x) for x in self.scales]
+        crop_h = [self.input_size[1] if abs(x - self.input_size[1]) < 3 else x for x in crop_sizes]
+        crop_w = [self.input_size[0] if abs(x - self.input_size[0]) < 3 else x for x in crop_sizes]
+
+        pairs = []
+        for i, h in enumerate(crop_h):
+            for j, w in enumerate(crop_w):
+                if abs(i - j) <= self.max_distort:
+                    pairs.append((w, h))
+
+        crop_pair = random.choice(pairs)
+        if not self.fix_crop:
+            w_offset = random.randint(0, image_w - crop_pair[0])
+            h_offset = random.randint(0, image_h - crop_pair[1])
+        else:
+            w_offset, h_offset = self._sample_fix_offset(image_w, image_h, crop_pair[0], crop_pair[1])
+
+        return crop_pair[0], crop_pair[1], w_offset, h_offset
+
+    def _sample_fix_offset(self, image_w, image_h, crop_w, crop_h):
+        offsets = self.fill_fix_offset(self.more_fix_crop, image_w, image_h, crop_w, crop_h)
+        return random.choice(offsets)
+
+    @staticmethod
+    def fill_fix_offset(more_fix_crop, image_w, image_h, crop_w, crop_h):
+        w_step = (image_w - crop_w) // 4
+        h_step = (image_h - crop_h) // 4
+
+        ret = list()
+        ret.append((0, 0))  # upper left
+        ret.append((4 * w_step, 0))  # upper right
+        ret.append((0, 4 * h_step))  # lower left
+        ret.append((4 * w_step, 4 * h_step))  # lower right
+        ret.append((2 * w_step, 2 * h_step))  # center
+
+        if more_fix_crop:
+            ret.append((0, 2 * h_step))  # center left
+            ret.append((4 * w_step, 2 * h_step))  # center right
+            ret.append((2 * w_step, 4 * h_step))  # lower center
+            ret.append((2 * w_step, 0 * h_step))  # upper center
+
+            ret.append((1 * w_step, 1 * h_step))  # upper left quarter
+            ret.append((3 * w_step, 1 * h_step))  # upper right quarter
+            ret.append((1 * w_step, 3 * h_step))  # lower left quarter
+            ret.append((3 * w_step, 3 * h_step))  # lower righ quarter
+
+        return ret
+
+
+class GroupRandomSizedCrop(object):
+    """Random crop the given PIL.Image to a random size of (0.08 to 1.0) of the original size
+    and and a random aspect ratio of 3/4 to 4/3 of the original aspect ratio
+    This is popularly used to train the Inception networks
+    size: size of the smaller edge
+    interpolation: Default: PIL.Image.BILINEAR
+    """
+    def __init__(self, size, interpolation=Image.BILINEAR):
+        self.size = size
+        self.interpolation = interpolation
+
+    def __call__(self, img_group):
+        for attempt in range(10):
+            area = img_group[0].size(0) * img_group[0].size(1)
+            target_area = random.uniform(0.08, 1.0) * area
+            aspect_ratio = random.uniform(3. / 4, 4. / 3)
+
+            w = int(round(math.sqrt(target_area * aspect_ratio)))
+            h = int(round(math.sqrt(target_area / aspect_ratio)))
+
+            if random.random() < 0.5:
+                w, h = h, w
+
+            if w <= img_group[0].size(0) and h <= img_group[0].size(1):
+                x1 = random.randint(0, img_group[0].size(0) - w)
+                y1 = random.randint(0, img_group[0].size(1) - h)
+                found = True
+                break
+        else:
+            found = False
+            x1 = 0
+            y1 = 0
+
+        if found:
+            out_group = list()
+            for img in img_group:
+                img = img.crop((x1, y1, x1 + w, y1 + h))
+                assert(img.size == (w, h))
+                out_group.append(img.resize((self.size, self.size), self.interpolation))
+            return out_group
+        else:
+            # Fallback
+            scale = GroupScale(self.size, interpolation=self.interpolation)
+            crop = GroupRandomCrop(self.size)
+            return crop(scale(img_group))
+
+
+class Stack(object):
+
+    def __init__(self, roll=False):
+        self.roll = roll
+
+    def __call__(self, img_group):
+        if img_group[0].mode == 'L':
+            return np.concatenate([np.expand_dims(x, 2) for x in img_group], axis=2)
+        elif img_group[0].mode == 'RGB':
+            if self.roll:
+                return np.concatenate([np.array(x)[:, :, ::-1] for x in img_group], axis=2)
+            else:
+                return np.concatenate(img_group, axis=2)
+
+# LSTR                
+class GroupStack(object):
+    """Stack to numpy array and convert  THWC to TCHW"""
+
+    def __call__(self, img_group):
+        """
+        Args:
+            img_group (List(PIL.Image)): a list of image. image shape: (H,W,C) for rgb, (H,W) for grayscale
+        Returns: np.ndarray. stack the images along the first axis.
+                shape: (T, C, H, W).
+        """
+        if isinstance(img_group, (list, tuple)):
+            imgs = np.stack(img_group, axis=0)
+        else:
+            imgs = img_group
+        if imgs.ndim == 3:    # THW
+            imgs = imgs[:, np.newaxis, :, :]    # TCHW
+        else:    # THWC
+            imgs = np.transpose(imgs, [0, 3, 1, 2])    # THWC -> TCHW
+        return imgs
+
+class ToTorchFormatTensor(object):
+    """ Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255]
+    to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] """
+    def __init__(self, div=True):
+        self.div = div
+
+    def __call__(self, pic):
+        if isinstance(pic, np.ndarray):
+            # handle numpy array
+            img = torch.from_numpy(pic).permute(2, 0, 1).contiguous()
+        else:
+            # handle PIL Image
+            img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
+            img = img.view(pic.size[1], pic.size[0], len(pic.mode))
+            # put it from HWC to CHW format
+            # yikes, this transpose takes 80% of the loading time/CPU
+            img = img.transpose(0, 1).transpose(0, 2).contiguous()
+        return img.float().div(255) if self.div else img.float()
+
+# LSTR
+class GroupToTensor(object):
+    
+    def __init__(self, div=True):
+        self.div = div
+
+    def __call__(self, imgs):
+        """
+        Args:
+            imgs (np.ndarray) : The stacked images. shape: (T, C, H, W). dtype: np.uint8
+        Returns: (torch.Tensor). convert dtype to torch.float32.
+                If `div` is True, the images are divided by 255 to be normalized to [0,1]
+            
+        """
+        imgs = torch.from_numpy(imgs).float()
+        if self.div:
+            imgs = imgs.div(255.)
+        return imgs
+
+class IdentityTransform(object):
+
+    def __call__(self, data):
+        return data
+
+
+if __name__ == "__main__":
+    trans = torchvision.transforms.Compose([
+        GroupScale(256),
+        GroupRandomCrop(224),
+        Stack(),
+        ToTorchFormatTensor(),
+        GroupNormalize(
+            mean=[.485, .456, .406],
+            std=[.229, .224, .225]
+        )]
+    )
+
+    im = Image.open('../tensorflow-model-zoo.torch/lena_299.png')
+
+    color_group = [im] * 3
+    rst = trans(color_group)
+
+    gray_group = [im.convert('L')] * 9
+    gray_rst = trans(gray_group)
+
+    trans2 = torchvision.transforms.Compose([
+        GroupRandomSizedCrop(256),
+        Stack(),
+        ToTorchFormatTensor(),
+        GroupNormalize(
+            mean=[.485, .456, .406],
+            std=[.229, .224, .225])
+    ])
+    print(trans2(color_group))

third_party/EgoSpeak/MiniROAD/utils/logger.py

@@ -0,0 +1,17 @@
+import logging
+import os
+
+def get_logger(output_path):
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.DEBUG) 
+    # formatter = logging.Formatter('%(asctime)s:%(message)s ', '%Y-%m-%d %H:%M:%S')
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+    # console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    log_file = os.path.join(output_path, 'log.txt')
+    fh = logging.FileHandler(log_file)
+    fh.setLevel(logging.INFO)
+    logger.addHandler(fh)
+
+    return logger

third_party/EgoSpeak/MiniROAD/utils/lr_scheduler.py

@@ -0,0 +1,90 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# code from detectron2
+
+import math
+from typing import List
+import torch
+
+try:
+    from torch.optim.lr_scheduler import LRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+
+class WarmupCosineLR(LRScheduler):
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        max_iters: int,
+        warmup_factor: float = 0.001,
+        warmup_iters: int = 1000,
+        warmup_method: str = "linear",
+        last_epoch: int = -1,
+    ):
+        self.max_iters = max_iters
+        self.warmup_factor = warmup_factor
+        self.warmup_iters = warmup_iters
+        self.warmup_method = warmup_method
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        warmup_factor = _get_warmup_factor_at_iter(
+            self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor
+        )
+        # Different definitions of half-cosine with warmup are possible. For
+        # simplicity we multiply the standard half-cosine schedule by the warmup
+        # factor. An alternative is to start the period of the cosine at warmup_iters
+        # instead of at 0. In the case that warmup_iters << max_iters the two are
+        # very close to each other.
+        return [
+            base_lr
+            * warmup_factor
+            * 0.5
+            * (1.0 + math.cos(math.pi * self.last_epoch / self.max_iters))
+            for base_lr in self.base_lrs
+        ]
+
+    def _compute_values(self) -> List[float]:
+        # The new interface
+        return self.get_lr()
+
+
+def _get_warmup_factor_at_iter(
+    method: str, iter: int, warmup_iters: int, warmup_factor: float
+) -> float:
+    """
+    Return the learning rate warmup factor at a specific iteration.
+    See :paper:`ImageNet in 1h` for more details.
+
+    Args:
+        method (str): warmup method; either "constant" or "linear".
+        iter (int): iteration at which to calculate the warmup factor.
+        warmup_iters (int): the number of warmup iterations.
+        warmup_factor (float): the base warmup factor (the meaning changes according
+            to the method used).
+
+    Returns:
+        float: the effective warmup factor at the given iteration.
+    """
+    if iter >= warmup_iters:
+        return 1.0
+
+    if method == "constant":
+        return warmup_factor
+    elif method == "linear":
+        alpha = iter / warmup_iters
+        return warmup_factor * (1 - alpha) + alpha
+    else:
+        raise ValueError("Unknown warmup method: {}".format(method))
+
+def build_lr_scheduler(cfg, optimizer, iters_per_epoch):
+    max_iters = cfg['num_epoch'] * iters_per_epoch
+    warmup_iters = int(cfg['warmup_epoch'] * iters_per_epoch)
+    # warmup_iters = int(max_iters * 1/6)
+    scheduler = WarmupCosineLR(
+        optimizer,
+        max_iters=max_iters,
+        warmup_factor=cfg['warmup_factor'],
+        warmup_iters=warmup_iters,
+        warmup_method=cfg['warmup_method'],
+        last_epoch=-1)
+    return scheduler

third_party/EgoSpeak/MiniROAD/utils/metrics.py

@@ -0,0 +1,226 @@
+# Evaluation code from LSTR
+
+from collections import OrderedDict
+from multiprocessing import Pool
+
+import numpy as np
+from sklearn.metrics import average_precision_score
+
+
+def calibrated_average_precision_score(y_true, y_score):
+    """Compute calibrated average precision (cAP), which is particularly
+    proposed for the TVSeries dataset.
+    """
+    y_true_sorted = y_true[np.argsort(-y_score)]
+    tp = y_true_sorted.astype(float)
+    fp = np.abs(y_true_sorted.astype(float) - 1)
+    tps = np.cumsum(tp)
+    fps = np.cumsum(fp)
+    ratio = np.sum(tp == 0) / np.sum(tp)
+    cprec = tps / (tps + fps / (ratio + np.finfo(float).eps) + np.finfo(float).eps)
+    cap = np.sum(cprec[tp == 1]) / np.sum(tp)
+    return cap
+
+def perframe_average_precision(prediction, ground_truth, class_names, ignore_index,
+                               postprocessing=None, metrics='AP'):
+    """Compute (frame-level) average precision between ground truth and
+    predictions data frames.
+    """
+    result = OrderedDict()
+    ground_truth = np.array(ground_truth)
+    prediction = np.array(prediction)
+
+    # Postprocessing
+    if postprocessing is not None:
+        ground_truth, prediction = postprocessing(ground_truth, prediction)
+        
+    # Build metrics
+    if metrics == 'AP':
+        compute_score = average_precision_score
+    elif metrics == 'cAP':
+        # print('cAP')
+        compute_score = calibrated_average_precision_score
+    else:
+        raise RuntimeError('Unknown metrics: {}'.format(metrics))
+
+    # Ignore backgroud class
+    ignore_set = set([ignore_index])
+    
+    # print('ignore_set = set([{ignore_index}])')
+    
+    # Enable background according to ignore_index
+    # ignore_set = set([ignore_index])
+    
+    # Compute average precision
+    result['per_class_AP'] = OrderedDict()
+    result['num'] = OrderedDict()
+    print(f"NUM FRAMES: {np.sum(ground_truth[:, 1:])}")
+    for idx, class_name in enumerate(class_names):
+        if idx not in ignore_set:
+            # ground_truth.shape [2901376, 2]
+            # prediction.shape [45334 ?]
+            if np.any(ground_truth[:, idx]):
+                ap_score = compute_score(ground_truth[:, idx], prediction[:, idx])
+                result['per_class_AP'][class_name] = ap_score
+                result['num'][class_name] = f'[true: {int(np.sum(ground_truth[:, idx]))}, pred:{int(np.sum(prediction[:,idx]))}, AP:{ap_score*100:.1f}]'
+    result['mean_AP'] = np.mean(list(result['per_class_AP'].values()))
+    print(f"Per_class_AP: {result['per_class_AP']}")
+    return result
+
+def event_accuracy(pred_scores_dict, gt_targets_dict):
+    th_session_accuracies = {}
+    
+    threshold = 0.5
+    # ERROR: something is wrong
+    # PLEASE EDIT:
+    # pred_scores_dict here has wrong dimension
+    # [64, 128, 10, 3] <- this is because I am using trainloader (do not evaluate on trainloader)
+    for threshold in np.arange(0.1, 1.0, 0.2): # 0.1, 0.3, 0.5, 0.7, 0.9
+        num_events = 0
+        th_total_speaking_start_correct = 0
+        th_total_speaking_end_correct = 0
+        
+        for session in pred_scores_dict.keys():       
+            th_session_speaking_start_correct = 0
+            th_session_speaking_start_predictions = 0
+            
+            th_session_speaking_end_correct = 0
+            th_session_speaking_end_predictions = 0
+            
+            argmax_session_speaking_start_correct = 0
+            argmax_session_speaking_end_correct = 0
+            
+            pred_scores = pred_scores_dict[session]
+            # pred_scores: [num_feature_frames, anticipation_steps, num_classes]
+            # classes: background, target speaker speaking, other speaker speaking
+            
+            # 12/06: To-do
+            # I forgot whether miniroad AA step 0 means anticipate 1 step future or current
+            # I did change something it would likely to be 
+            ant_length = pred_scores.shape[1]
+            
+            gt_targets = gt_targets_dict[session]
+            event_labels = np.load(f'/home/junhyeok/projects/turn-taking/datasets/EasyCom/target_perframe/events/{session}.npy')
+            
+            num_speaking_start = 0
+            num_speaking_end = 0
+            print(f'{session} loaded, {len(event_labels)} events')
+            # 12/06:
+            # IndexError: index 8504 is out of bounds for axis 0 with size 8390
+            # why the event_idx has much higher value than the length of pred_scores?
+            # check how the event label is made  
+            
+            # event_label: [event_idx, event_class]
+            # ['18', 'target speaker speaking start']
+            for event_label in event_labels:
+                event_idx, event_class = event_label
+                event_idx = int(event_idx)
+                
+                if event_class == 'target speaker speaking start':
+                    num_speaking_start += 1
+                    num_speaking_end += 1
+                    # right before the event start, want to see how the model predict?
+                    ant_prob = pred_scores[event_idx-1, 0, :]
+                    if np.argmax(ant_prob) == 1:
+                        argmax_session_speaking_start_correct += 1
+                    
+                    ant_target_prob = pred_scores[event_idx-1, 0, 1]
+                    if ant_target_prob >= threshold:
+                        th_session_speaking_start_correct += 1
+                    
+                elif event_class == 'target speaker speaking end':
+                    # right after the event end, want to see how the model predict?
+                    ant_prob = pred_scores[event_idx, 0, :]
+                    if np.argmax(ant_prob) == 0:
+                        argmax_session_speaking_end_correct += 1
+                    
+                    # target speaker will not speaker anymore
+                    ant_target_prob = pred_scores[event_idx, 0, 1]
+                    if ant_target_prob < threshold:
+                        th_session_speaking_end_correct += 1
+            
+            print(f'Threshold {threshold:.2f}) {session} speaking start accuracy: {(th_session_speaking_start_correct / (len(event_labels) / 2)):.4f}')
+            print(f'Threshold {threshold:.2f}) {session} speaking end accuracy: {(th_session_speaking_end_correct / (len(event_labels)/ 2 )):.4f}')
+            print(f'argmax {session} speaking start accuracy: {argmax_session_speaking_start_correct / (len(event_labels) / 2)}')
+            print(f'argmax {session} speaking end accuracy: {argmax_session_speaking_end_correct / (len(event_labels)/ 2 )}')
+            th_session_accuracies[session] = [th_session_speaking_start_correct / len(event_labels), th_session_speaking_end_correct / len(event_labels)]
+            th_total_speaking_start_correct += th_session_speaking_start_correct
+            th_total_speaking_end_correct += th_session_speaking_end_correct
+            num_events += len(event_labels)
+            
+        th_total_speaking_start_accuracy = th_total_speaking_start_correct / len(event_labels)
+        th_total_speaking_end_accuracy = th_total_speaking_end_correct / len(event_labels)
+        
+        th_total_accuracy = (th_total_speaking_start_correct + th_total_speaking_end_correct) / (num_events)
+        print(f'total speaking start accuracy: {th_total_accuracy}')
+        
+    return th_total_accuracy
+    
+def get_stage_pred_scores(gt_targets, pred_scores, perc_s, perc_e):
+    starts = []
+    ends = []
+    stage_gt_targets = []
+    stage_pred_scores = []
+    for i in range(len(gt_targets)):
+        if gt_targets[i] == 0:
+            stage_gt_targets.append(gt_targets[i])
+            stage_pred_scores.append(pred_scores[i])
+        else:
+            if i == 0 or gt_targets[i - 1] == 0:
+                starts.append(i)
+            if i == len(gt_targets) - 1 or gt_targets[i + 1] == 0:
+                ends.append(i)
+    if len(starts) != len(ends):
+        raise ValueError('starts and ends cannot pair!')
+
+    action_lens = [ends[i] - starts[i] for i in range(len(starts))]
+    stage_starts = [starts[i] + int(action_lens[i] * perc_s) for i in range(len(starts))]
+    stage_ends = [max(stage_starts[i] + 1, starts[i] + int(action_lens[i] * perc_e)) for i in range(len(starts))]
+    for i in range(len(starts)):
+        stage_gt_targets.extend(gt_targets[stage_starts[i]: stage_ends[i]])
+        stage_pred_scores.extend(pred_scores[stage_starts[i]: stage_ends[i]])
+    return np.array(stage_gt_targets), np.array(stage_pred_scores)
+
+
+def perstage_average_precision(prediction, ground_truth,
+                               class_names, postprocessing, 
+                               metrics='cAP'):
+    result = OrderedDict()
+    ground_truth = np.array(ground_truth)
+    prediction = np.array(prediction)
+
+    # Postprocessing
+    if postprocessing is not None:
+        ground_truth, prediction = postprocessing(ground_truth, prediction)
+
+    # Build metrics
+    if metrics == 'AP':
+        compute_score = average_precision_score
+    elif metrics == 'cAP':
+        compute_score = calibrated_average_precision_score
+    else:
+        raise RuntimeError('Unknown metrics: {}'.format(metrics))
+
+    raise NotImplementedError('This function is not implemented yet.')
+    # Ignore backgroud class
+    # ignore_index = set([0])
+
+    # Compute average precision
+    for perc_s in range(10):
+        perc_e = perc_s + 1
+        stage_name = '{:2}%_{:3}%'.format(perc_s * 10, perc_e * 10)
+        result[stage_name] = OrderedDict({'per_class_AP': OrderedDict()})
+        for idx, class_name in enumerate(class_names):
+            if idx not in ignore_index:
+                stage_gt_targets, stage_pred_scores = get_stage_pred_scores(
+                    (ground_truth[:, idx] == 1).astype(int),
+                    prediction[:, idx],
+                    perc_s / 10,
+                    perc_e / 10,
+                )
+                result[stage_name]['per_class_AP'][class_name] = \
+                    compute_score(stage_gt_targets, stage_pred_scores)
+        result[stage_name]['mean_AP'] = \
+            np.mean(list(result[stage_name]['per_class_AP'].values()))
+
+    return result

third_party/EgoSpeak/MiniROAD/utils/postprocessing.py

@@ -0,0 +1,29 @@
+
+import numpy as np
+
+def thumos_postprocessing(ground_truth, prediction, smooth=False, switch=False):
+    """
+    We follow (Shou et al., 2017) and adopt their perframe postprocessing method on THUMOS'14 datset.
+    Source: https://bitbucket.org/columbiadvmm/cdc/src/master/THUMOS14/eval/PreFrameLabeling/compute_framelevel_mAP.m
+    """
+
+    # Simple temporal smoothing via NMS of 5-frames window
+    if smooth:
+        prob = np.copy(prediction)
+        prob1 = prob.reshape(1, prob.shape[0], prob.shape[1])
+        prob2 = np.append(prob[0, :].reshape(1, -1), prob[0: -1, :], axis=0).reshape(1, prob.shape[0], prob.shape[1])
+        prob3 = np.append(prob[1:, :], prob[-1, :].reshape(1, -1), axis=0).reshape(1, prob.shape[0], prob.shape[1])
+        prob4 = np.append(prob[0: 2, :], prob[0: -2, :], axis=0).reshape(1, prob.shape[0], prob.shape[1])
+        prob5 = np.append(prob[2:, :], prob[-2:, :], axis=0).reshape(1, prob.shape[0], prob.shape[1])
+        probsmooth = np.squeeze(np.max(np.concatenate((prob1, prob2, prob3, prob4, prob5), axis=0), axis=0))
+        prediction = np.copy(probsmooth)
+
+    # Assign cliff diving (5) as diving (8)
+    if switch:
+        switch_index = np.where(prediction[:, 5] > prediction[:, 8])[0]
+        prediction[switch_index, 8] = prediction[switch_index, 5]
+
+    # Remove ambiguous (21)
+    valid_index = np.where(ground_truth[:, 21] != 1)[0]
+
+    return ground_truth[valid_index], prediction[valid_index]

third_party/EgoSpeak/MiniROAD/utils/registry.py

@@ -0,0 +1,20 @@
+def _register_generic(module_dict, module_name, module):
+    assert module_name not in module_dict
+    module_dict[module_name] = module
+
+
+class Registry(dict):
+
+    def __init__(self, *args, **kwargs):
+        super(Registry, self).__init__(*args, **kwargs)
+
+    def register(self, module_name, module=None):
+        if module is not None:
+            _register_generic(self, module_name, module)
+            return
+
+        def register_fn(fn):
+            _register_generic(self, module_name, fn)
+            return fn
+
+        return register_fn

third_party/EgoSpeak/MiniROAD/utils/util.py

@@ -0,0 +1,74 @@
+import pickle
+import os
+import os.path as osp
+import random
+import numpy as np
+import torch
+
+def dump_pickle(lst, file_path, file_name):
+    with open(osp.join(file_path, file_name + '.pkl'), 'wb') as f:
+        pickle.dump(lst, f)
+
+def create_dir(dir_path):
+    if not osp.exists(dir_path):
+        os.makedirs(dir_path)
+
+def create_outdir(result_path, overwrite=False):
+    i = 1
+    new_result_path = result_path
+    
+    if overwrite == False:
+        while osp.exists(new_result_path):
+            new_result_path = f'{result_path}_{i}'
+            i += 1
+    
+    create_dir(osp.join(new_result_path, 'ckpts'))
+    create_dir(osp.join(new_result_path, 'runs'))
+    return new_result_path
+
+def set_seed(seed):
+    # os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    
+def parse_rgb_type(rgb_type):    
+    if "rgb_kinetics_resnet50" in rgb_type:
+        rgb_type = "rgb_kinetics_resnet50"
+    elif "rgb_farl" in rgb_type:
+        rgb_type = "rgb_farl"
+    
+    return rgb_type
+
+def parse_flow_type(flow_type):
+    if "audio_wav2vec2" in flow_type:
+        flow_type = "audio_wav2vec2"
+    
+    return flow_type
+
+def get_identifier(cfg):
+    mode = "eval" if cfg["eval"] == True else "train"
+    
+    modality = ''
+    if not cfg["no_flow"]:
+        modality += 'A'
+    if not cfg["no_rgb"]:
+        modality += 'V'
+    
+    identifier = f'{cfg["model"]}_{cfg["data_name"]}'
+    if 'A' in modality:
+        identifier += f"_{cfg['flow_type']}"
+    if 'V' in modality:
+        identifier += f"_{cfg['rgb_type']}"
+    
+    identifier += f'_{modality}_epoch{cfg["num_epoch"]}_{mode}'
+    if cfg["ckpt"] != None:
+        identifier += f'_ckpt_{".".join(osp.basename(cfg["ckpt"]).split(".")[:-1])}'
+    
+    return identifier
+