Add causal person fusion and MuVAP stage-3 baselines

third_party/EgoSpeak/LSTR/src/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/__init__.py

@@ -0,0 +1,4 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .defaults import get_cfg

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/defaults.py

@@ -0,0 +1,171 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from yacs.config import CfgNode as CN
+
+# ---------------------------------------------------------------------------- #
+# Config Definition
+# ---------------------------------------------------------------------------- #
+_C = CN()
+
+# ---------------------------------------------------------------------------- #
+# Metadata
+# ---------------------------------------------------------------------------- #
+_C.SEED = 0
+
+# ---------------------------------------------------------------------------- #
+# Model
+# ---------------------------------------------------------------------------- #
+_C.MODEL = CN()
+_C.MODEL.MODEL_NAME = ''
+_C.MODEL.CHECKPOINT = ''
+_C.MODEL.EXCEPT_CLASSIFIER = False
+
+# ---------------------------------------------------------------------------- #
+# Feature Head
+# ---------------------------------------------------------------------------- #
+_C.MODEL.FEATURE_HEAD = CN()
+_C.MODEL.FEATURE_HEAD.LINEAR_ENABLED = True
+_C.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES = 1024
+_C.MODEL.FEATURE_HEAD.DROPOUT = 0.
+
+# ---------------------------------------------------------------------------- #
+# Transformer Network
+# ---------------------------------------------------------------------------- #
+_C.MODEL.LSTR = CN()
+# Hyperparameters
+_C.MODEL.LSTR.NUM_HEADS = 8
+_C.MODEL.LSTR.DIM_FEEDFORWARD = 1024
+_C.MODEL.LSTR.DROPOUT = 0.2
+_C.MODEL.LSTR.DROPOUT_CLS = 0.0
+_C.MODEL.LSTR.ACTIVATION = 'relu'
+# Memory choices
+_C.MODEL.LSTR.AGES_MEMORY_SECONDS = 0
+_C.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE = 1
+_C.MODEL.LSTR.LONG_MEMORY_SECONDS = 0
+_C.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE = 1
+_C.MODEL.LSTR.WORK_MEMORY_SECONDS = 8
+_C.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE = 1
+_C.MODEL.LSTR.WORK_MEMORY_USE_PE = True
+_C.MODEL.LSTR.LONG_MEMORY_USE_PE = False
+_C.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK = False
+_C.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2 = False
+# Anticipation choices
+_C.MODEL.LSTR.ANTICIPATION_SECONDS = 0
+_C.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE = 1
+# for EK55/100
+_C.MODEL.LSTR.V_N_CLASSIFIER = False
+_C.MODEL.LSTR.LOSS_ANTICIPATE_ONLY = False
+_C.MODEL.LSTR.LOSS_WEIGHT_V = 1.
+_C.MODEL.LSTR.LOSS_WEIGHT_N = 1.
+_C.MODEL.LSTR.FUSION_WEIGHT_V = 1.
+_C.MODEL.LSTR.FUSION_WEIGHT_N = 1.
+# Design choices
+_C.MODEL.LSTR.ENC_ATTENTION_TYPE = 'dotproduct'
+_C.MODEL.LSTR.ENC_ATTENTION_DECAY = 1.0
+_C.MODEL.LSTR.ENC_MODULE = [
+    [16, 1, True], [32, 2, True]
+]
+_C.MODEL.LSTR.DEC_MODULE = [-1, 2, True]
+_C.MODEL.LSTR.FC_NORM = False
+# Inference modes
+_C.MODEL.LSTR.INFERENCE_MODE = 'batch'
+# Position Embedding
+# _C.MODEL.LSTR.POSITION_EMBEDDING = 'default'
+
+# ---------------------------------------------------------------------------- #
+# Criterion
+# ---------------------------------------------------------------------------- #
+_C.MODEL.CRITERIONS = [['MCE', {}]]
+
+# ---------------------------------------------------------------------------- #
+# Data
+# ---------------------------------------------------------------------------- #
+_C.DATA = CN()
+_C.DATA.DATA_INFO = 'data/data_info.json'
+_C.DATA.DATA_NAME = None
+_C.DATA.DATA_ROOT = None
+_C.DATA.CLASS_NAMES = None
+_C.DATA.NUM_CLASSES = None
+_C.DATA.IGNORE_INDEX = None
+_C.DATA.METRICS = None
+_C.DATA.FPS = None
+_C.DATA.TRAIN_SESSION_SET = None
+_C.DATA.TEST_SESSION_SET = None
+_C.DATA.EK_EXT_PATH = 'external/rulstm/RULSTM/data/ek55/'
+_C.DATA.CLIP_MIXUP_RATE = 0.
+_C.DATA.CLIP_MIXUP_SAMPLE = 'uniform'
+
+# ---------------------------------------------------------------------------- #
+# Input
+# ---------------------------------------------------------------------------- #
+_C.INPUT = CN()
+_C.INPUT.MODALITY = 'twostream'
+# 'twostream'='visual+motion', 'threestream'='visual+motion+object'
+_C.INPUT.VISUAL_FEATURE = 'rgb_anet_resnet50'
+_C.INPUT.MOTION_FEATURE = 'flow_anet_resnet50'
+_C.INPUT.OBJECT_FEATURE = 'flow_anet_resnet50'
+_C.INPUT.TARGET_PERFRAME = 'target_perframe'
+
+# ---------------------------------------------------------------------------- #
+# Data Loader
+# ---------------------------------------------------------------------------- #
+_C.DATA_LOADER = CN()
+_C.DATA_LOADER.BATCH_SIZE = 32
+_C.DATA_LOADER.NUM_WORKERS = 4
+_C.DATA_LOADER.PIN_MEMORY = False
+
+# _C.MODEL.MINIROAD.NUM_WORKERS = 4
+
+# ---------------------------------------------------------------------------- #
+# Solver
+# ---------------------------------------------------------------------------- #
+_C.SOLVER = CN()
+_C.SOLVER.START_EPOCH = 1
+_C.SOLVER.NUM_EPOCHS = 50
+
+# ---------------------------------------------------------------------------- #
+# Optimizer
+# ---------------------------------------------------------------------------- #
+_C.SOLVER.OPTIMIZER = 'adam'
+_C.SOLVER.BASE_LR = 0.00005
+_C.SOLVER.WEIGHT_DECAY = 0.00005
+_C.SOLVER.MOMENTUM = 0.9
+
+# _C.MODEL.MINIROAD.LR = 0.0001
+# _C.MODEL.MINIROAD.WEIGHT_DECAY = 0.05
+
+
+# ---------------------------------------------------------------------------- #
+# Scheduler
+# ---------------------------------------------------------------------------- #
+_C.SOLVER.SCHEDULER = CN()
+_C.SOLVER.SCHEDULER.SCHEDULER_NAME = None
+_C.SOLVER.SCHEDULER.MILESTONES = []
+_C.SOLVER.SCHEDULER.GAMMA = 0.1
+_C.SOLVER.SCHEDULER.WARMUP_FACTOR = 0.3
+_C.SOLVER.SCHEDULER.WARMUP_EPOCHS = 10.0
+_C.SOLVER.SCHEDULER.WARMUP_METHOD = 'linear'
+
+# ---------------------------------------------------------------------------- #
+# Others
+# ---------------------------------------------------------------------------- #
+_C.SOLVER.PHASES = ['train', 'test']
+_C.SOLVER.SAVE_EVERY_EPOCHS = 10
+
+# ---------------------------------------------------------------------------- #
+# Output
+# ---------------------------------------------------------------------------- #
+_C.OUTPUT_DIR = 'checkpoints'
+_C.SESSION = ''
+_C.USE_WANDB = False
+_C.SAVE_HPARAM = False
+
+# ---------------------------------------------------------------------------- #
+# Misc
+# ---------------------------------------------------------------------------- #
+_C.VERBOSE = False
+
+
+def get_cfg():
+    return _C.clone()

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/feature_sizes.py

@@ -0,0 +1,20 @@
+# 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,
+}

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/__init__.py

@@ -0,0 +1,4 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .criterions import build_criterion

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/criterions.py

@@ -0,0 +1,356 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['build_criterion']
+
+import os.path as osp
+from functools import partial
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rekognition_online_action_detection.utils.registry import Registry
+
+CRITERIONS = Registry()
+
+
+@CRITERIONS.register('BCE')
+class BinaryCrossEntropyLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(BinaryCrossEntropyLoss, self).__init__()
+
+        self.criterion = nn.BCEWithLogitsLoss(reduction=reduction)
+
+    def forward(self, input, target):
+        return self.criterion(input, target)
+
+
+@CRITERIONS.register('SCE')
+class SingleCrossEntropyLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(SingleCrossEntropyLoss, self).__init__()
+
+        self.criterion = nn.CrossEntropyLoss(
+            reduction=reduction, ignore_index=ignore_index)
+
+    def forward(self, input, target):
+        return self.criterion(input, target)
+
+
+@CRITERIONS.register('MCE')
+class MultipCrossEntropyLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(MultipCrossEntropyLoss, self).__init__()
+
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+
+    def forward(self, input, target):
+        logsoftmax = nn.LogSoftmax(dim=-1).to(input.device)
+
+        # LSTR transform input from [16,50,3] to [800,3]
+        output = torch.sum(-target * logsoftmax(input), dim=1)
+
+        if self.reduction == 'mean':
+            return torch.mean(output)
+        elif self.reduction == 'sum':
+            return torch.sum(output)
+        else:
+            return output
+        
+        # if self.ignore_index >= 0:
+        #     notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index]
+        #     output = torch.sum(-target[:, notice_index] * logsoftmax(input[:, notice_index]), dim=1)
+        #     if (target[:, self.ignore_index] != 1).sum().item() == 0:
+        #         return torch.mean(torch.zeros_like(output))
+        #     if self.reduction == 'mean':
+        #         return torch.mean(output[target[:, self.ignore_index] != 1])
+        #     elif self.reduction == 'sum':
+        #         return torch.sum(output[target[:, self.ignore_index] != 1])
+        #     else:
+        #         return output[target[:, self.ignore_index] != 1]
+        # else:
+        #     output = torch.sum(-target * logsoftmax(input), dim=1)
+
+        #     if self.reduction == 'mean':
+        #         return torch.mean(output)
+        #     elif self.reduction == 'sum':
+        #         return torch.sum(output)
+        #     else:
+        #         return output
+
+@CRITERIONS.register('MCE_MiniROAD')
+class MiniROADMultipCrossEntropyLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(MiniROADMultipCrossEntropyLoss, self).__init__()
+
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+
+    def forward(self, input, target):
+        logsoftmax = nn.LogSoftmax(dim=-1).to(input.device)
+        output = torch.sum(-F.normalize(target) * logsoftmax(input), dim=1) # B
+        if self.reduction == 'mean':
+            loss = torch.mean(output)
+        elif self.reduction == 'sum':
+            loss = torch.sum(output)
+
+        return loss
+    
+@CRITERIONS.register('MCE_ANT')
+class AntMultipCrossEntropyLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(AntMultipCrossEntropyLoss, self).__init__()
+
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+
+    # only ant_target loss
+    def forward(self, input, target):
+        logsoftmax = nn.LogSoftmax(dim=-1).to(input.device)
+        import pdb; pdb.set_trace()
+        if self.ignore_index >= 0:
+            notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index]
+            import pdb; pdb.set_trace()
+            # input dim: [16, 128, 10, 3]
+            # target.shape: [16, 10, 3]
+            output = torch.sum()
+            
+            output = torch.sum(-target[:, notice_index] * logsoftmax(input[:, notice_index]), dim=1)
+            if (target[:, self.ignore_index] != 1).sum().item() == 0:
+                return torch.mean(torch.zeros_like(output))
+            if self.reduction == 'mean':
+                return torch.mean(output[target[:, self.ignore_index] != 1])
+            elif self.reduction == 'sum':
+                return torch.sum(output[target[:, self.ignore_index] != 1])
+            else:
+                return output[target[:, self.ignore_index] != 1]
+        else:
+            output = torch.sum(-target * logsoftmax(input), dim=1)
+
+            if self.reduction == 'mean':
+                return torch.mean(output)
+            elif self.reduction == 'sum':
+                return torch.sum(output)
+            else:
+                return output
+
+
+@CRITERIONS.register('MCE_EQL')
+class MultipCrossEntropyEqualizedLoss(nn.Module):
+
+    def __init__(self, gamma=0.95, lambda_=3e-3, reduction='mean', ignore_index=-100,
+                 anno_path='external/rulstm/RULSTM/data/ek55/'):
+        super(MultipCrossEntropyEqualizedLoss, self).__init__()
+
+        # get label distribution
+        segment_list = pd.read_csv(osp.join(anno_path, 'training.csv'),
+                                   names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+                                   skipinitialspace=True)
+        freq_info = np.zeros((max(segment_list['action']) + 1,))
+        assert ignore_index == 0
+        for segment in segment_list.iterrows():
+            freq_info[segment[1]['action']] += 1.
+        freq_info = freq_info / freq_info.sum()
+        self.freq_info = torch.FloatTensor(freq_info)
+
+        self.gamma = gamma
+        self.lambda_ = lambda_
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+
+    def forward(self, input, target):
+        logsoftmax = nn.LogSoftmax(dim=1).to(input.device)
+
+
+        bg_target = target[:, self.ignore_index]
+        notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index]
+        input = input[:, notice_index]
+        target = target[:, notice_index]
+
+        weight = input.new_zeros(len(notice_index))
+        weight[self.freq_info < self.lambda_] = 1.
+        weight = weight.view(1, -1).repeat(input.shape[0], 1)
+        
+        eql_w = 1 - (torch.rand_like(target) < self.gamma) * weight * (1 - target)
+        input = torch.log(eql_w + 1e-8) + input
+
+        output = torch.sum(-target * logsoftmax(input), dim=1)
+        if (bg_target != 1).sum().item() == 0:
+            return torch.mean(torch.zeros_like(output))
+        if self.reduction == 'mean':
+            return torch.mean(output[bg_target != 1])
+        elif self.reduction == 'sum':
+            return torch.sum(output[bg_target != 1])
+        else:
+            return output[bg_target != 1]
+
+
+@CRITERIONS.register('EQLv2')
+class EQLv2Loss(nn.Module):
+
+    def __init__(self, gamma=12, mu=0.8, alpha=4.0, reduction='mean', ignore_index=-100,
+                 num_classes=3806):
+        super(EQLv2Loss, self).__init__()
+        self.num_classes = num_classes
+
+        self.gamma = gamma
+        self.mu = mu
+        self.alpha = alpha
+        self.ignore_index = ignore_index
+
+        self._pos_grad = None
+        self._neg_grad = None
+        self.pos_neg = None
+
+        def _func(x, gamma, mu):
+            return 1 / (1 + torch.exp(-gamma * (x - mu)))
+        self.map_func = partial(_func, gamma=self.gamma, mu=self.mu)
+
+    def forward(self, input, target):
+        self.n_i, self.n_c = input.size()
+
+        pos_w, neg_w = self.get_weight(input)
+
+        weight = pos_w * target + neg_w * (1 - target)
+
+        cls_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
+        cls_loss = torch.sum(cls_loss * weight) / self.n_i
+
+        self.collect_grad(input.detach(), target.detach(), weight.detach())
+
+        return cls_loss
+
+    def collect_grad(self, cls_score, target, weight):
+        prob = torch.sigmoid(cls_score)
+        grad = target * (prob - 1) + (1 - target) * prob
+        grad = torch.abs(grad)
+
+        # do not collect grad for objectiveness branch [:-1]
+        pos_grad = torch.sum(grad * target * weight, dim=0)[1:]
+        neg_grad = torch.sum(grad * (1 - target) * weight, dim=0)[1:]
+
+        # dist.all_reduce(pos_grad)
+        # dist.all_reduce(neg_grad)
+
+        self._pos_grad += pos_grad
+        self._neg_grad += neg_grad
+        self.pos_neg = self._pos_grad / (self._neg_grad + 1e-10)
+
+    def get_weight(self, cls_score):
+        if self._pos_grad is None:
+            self._pos_grad = cls_score.new_zeros(self.num_classes)
+            self._neg_grad = cls_score.new_zeros(self.num_classes)
+            neg_w = cls_score.new_zeros((self.n_i, self.n_c))
+            pos_w = cls_score.new_zeros((self.n_i, self.n_c))
+        else:
+            neg_w = torch.cat([cls_score.new_ones(1), self.map_func(self.pos_neg)])
+            pos_w = 1 + self.alpha * (1 - neg_w)
+            neg_w = neg_w.view(1, -1).expand(self.n_i, self.n_c)
+            pos_w = pos_w.view(1, -1).expand(self.n_i, self.n_c)
+        return pos_w, neg_w
+
+
+@CRITERIONS.register('MCE_SEESAW')
+class MultipCrossEntropySeesawLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100,
+                 p=0.8, q=2.0, eps=1e-2,
+                 anno_path='external/rulstm/RULSTM/data/ek55/'):
+        super(MultipCrossEntropySeesawLoss, self).__init__()
+
+        # get label distribution
+        segment_list = pd.read_csv(osp.join(anno_path, 'training.csv'),
+                                   names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+                                   skipinitialspace=True)
+        self.num_classes = max(segment_list['action']) + 1
+        self.cum_samples = torch.zeros(self.num_classes, dtype=torch.float)
+
+        self.p = p
+        self.q = q
+        self.eps = eps
+
+        self.reduction = reduction
+        assert ignore_index == 0
+        self.ignore_index = ignore_index
+
+    def forward(self, input, target):
+        logsoftmax = nn.LogSoftmax(dim=1).to(input.device)
+        self.cum_samples = self.cum_samples.to(input.device)
+
+        bg_target = target[:, self.ignore_index]
+        notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index]
+        input = input[:, notice_index]
+        target = target[:, notice_index]
+
+        # accumulate the samples for each category
+        self.cum_samples += target.sum(0)
+        
+        seesaw_weights = input.new_ones(target.size())
+
+        if self.p > 0:
+            sample_ratio_matrix = self.cum_samples[None, :].clamp(min=1) / self.cum_samples[:, None].clamp(min=1)
+            index = (sample_ratio_matrix < 1.0).float()
+            sample_weights = sample_ratio_matrix.pow(self.p) * index + (1 - index)
+            for ni in range(input.shape[0]):
+                true_target = target[ni, :].nonzero(as_tuple=True)[0]
+                ## TODO: study mean/min/max/sum
+                if len(true_target) > 0:
+                    mitigation_factor = torch.min(sample_weights[true_target, :], dim=0).values
+                    seesaw_weights[ni, :] = seesaw_weights[ni, :] * mitigation_factor
+        if self.q > 0:
+            scores = F.softmax(input.detach(), dim=1)
+            self_scores = input.new_ones([input.shape[0]])
+            for ni in range(input.shape[0]):
+                true_target = target[ni, :].nonzero(as_tuple=True)[0]
+                ## TODO: study mean/min/max/sum
+                if len(true_target) > 0:
+                    self_scores[ni] = torch.min(scores[ni, true_target])
+            score_matrix = scores / self_scores[:, None].clamp(min=self.eps)
+            index = (score_matrix > 1.0).float()
+            compensation_factor = score_matrix.pow(self.q) * index + (1 - index)
+            seesaw_weights = seesaw_weights * compensation_factor
+
+        input = input + (seesaw_weights.log() * (1 - target))
+
+        output = torch.sum(-target * logsoftmax(input), dim=1)
+        if (bg_target != 1).sum().item() == 0:
+            return torch.mean(torch.zeros_like(output))
+        if self.reduction == 'mean':
+            return torch.mean(output[bg_target != 1])
+        elif self.reduction == 'sum':
+            return torch.sum(output[bg_target != 1])
+        else:
+            return output[bg_target != 1]
+
+
+@CRITERIONS.register('PRED_FUTURE')
+class PredictFutureLoss(nn.Module):
+
+    def __init__(self, reduction='mean', ignore_index=-100):
+        super(PredictFutureLoss, self).__init__()
+
+        self.reduction = reduction
+        self.criterion = nn.MSELoss(reduction=reduction)
+
+    def forward(self, pred, true):
+        return self.criterion(pred, true)
+
+
+def build_criterion(cfg, device=None):
+    criterion = {}
+    for name, params in cfg.MODEL.CRITERIONS:
+        if name in CRITERIONS:
+            if 'ignore_index' not in params:
+                params['ignore_index'] = cfg.DATA.IGNORE_INDEX
+            criterion[name] = CRITERIONS[name](**params).to(device)
+        else:
+            raise RuntimeError('Unknown criterion: {}'.format(name))
+    return criterion

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/__init__.py

@@ -0,0 +1,5 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0                                                                    
+
+from .datasets import build_data_loader, build_dataset
+from .perframe_data_layers import *

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/datasets.py

@@ -0,0 +1,31 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+                                     
+
+__all__ = [
+    'build_dataset',
+    'build_data_loader',
+]
+
+import torch.utils.data as data
+from rekognition_online_action_detection.utils.registry import Registry
+
+DATA_LAYERS = Registry()
+
+
+def build_dataset(cfg, phase, tag=''):
+    data_names = ['YTConv', 'EGO4D']
+    data_name = next((name for name in data_names if name in cfg.DATA.DATA_NAME), cfg.DATA.DATA_NAME)
+    
+    data_layer = DATA_LAYERS[cfg.MODEL.MODEL_NAME + tag + data_name]
+    return data_layer(cfg, phase)
+
+
+def build_data_loader(cfg, phase):
+    data_loader = data.DataLoader(
+        dataset=build_dataset(cfg, phase),
+        batch_size=cfg.DATA_LOADER.BATCH_SIZE,
+        shuffle=True if phase == 'train' else False,
+        num_workers=cfg.DATA_LOADER.NUM_WORKERS,
+        pin_memory=cfg.DATA_LOADER.PIN_MEMORY,
+    )
+    return data_loader

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/perframe_data_layers.py

@@ -0,0 +1,453 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0                                                                    
+                                     
+
+import os.path as osp
+import warnings
+from bisect import bisect_right
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.utils.data as data
+from rekognition_online_action_detection.utils.ek_utils import (
+    action_to_noun_map, action_to_verb_map)
+
+from .datasets import DATA_LAYERS as registry
+
+
+@registry.register('LSTRTHUMOS')
+@registry.register('LSTRTVSeries')
+@registry.register('LSTREK55')
+@registry.register('LSTREK100')
+@registry.register('LSTRYTConv')
+@registry.register('LSTREGO4D')
+@registry.register('LSTREasyCom')
+class LSTRDataLayer(data.Dataset):
+
+    def __init__(self, cfg, phase='train'):
+        self.cfg = cfg
+        self.data_name = cfg.DATA.DATA_NAME
+        self.data_root = cfg.DATA.DATA_ROOT
+        self.clip_mixup_rate = cfg.DATA.CLIP_MIXUP_RATE
+        self.clip_mixup_sample = cfg.DATA.CLIP_MIXUP_SAMPLE
+        self.visual_feature = cfg.INPUT.VISUAL_FEATURE
+        self.motion_feature = cfg.INPUT.MOTION_FEATURE
+        self.object_feature = cfg.INPUT.OBJECT_FEATURE
+        self.target_perframe = cfg.INPUT.TARGET_PERFRAME
+        self.sessions = getattr(cfg.DATA, phase.upper() + '_SESSION_SET')
+        self.long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH
+        self.long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE
+        self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES
+        self.work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH
+        self.work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE
+        self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES
+        self.anticipation_length = cfg.MODEL.LSTR.ANTICIPATION_LENGTH
+        self.anticipation_sample_rate = cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE
+        self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES
+        self.training = phase == 'train'
+        self._init_dataset()
+
+    def shuffle(self):
+        self._init_dataset()
+
+    def _init_dataset(self):
+        self.inputs = []
+        
+        if self.data_name.startswith('EK'):
+            if self.data_name == 'EK55':
+                path_to_data = 'external/rulstm/RULSTM/data/ek55/'
+            elif self.data_name == 'EK100':
+                path_to_data = 'external/rulstm/RULSTM/data/ek100/'
+            else:
+                raise ValueError
+            segment_list = pd.read_csv(osp.join(path_to_data,
+                                                'training.csv' if self.training else 'validation.csv'),
+                                       names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+                                       skipinitialspace=True)
+            segment_list['len_f'] = segment_list['end_f'] - segment_list['start_f']
+
+            self.segment_by_cls = {cls: segment_list[segment_list['action'] == cls]
+                                   for cls in range(0, self.cfg.DATA.NUM_CLASSES - 1)}
+            for session in self.sessions:
+                target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy'))
+                if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                    verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy'))
+                    noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy'))
+                else:
+                    verb_target = target
+                    noun_target = target
+                segments_per_session = segment_list[segment_list['video'] == session]
+                for segment in segments_per_session.iterrows():
+                    start_tick = int(segment[1]['start_f'] / 30 * self.cfg.DATA.FPS)
+                    end_tick = int(segment[1]['end_f'] / 30 * self.cfg.DATA.FPS)
+                    start_tick += np.random.randint(self.anticipation_length) if self.training else 0
+                    start_tick = min(start_tick, end_tick)
+                    work_end = start_tick - self.anticipation_length
+                    work_start = work_end - self.work_memory_length
+                    segments_before_current = segments_per_session[segments_per_session['end_f'] < segment[1]['start_f']]
+                    if work_start < 0:
+                        continue
+                    self.inputs.append([
+                        session, work_start, work_end,
+                        target[work_start: work_end + self.anticipation_length],
+                        verb_target[work_start: work_end + self.anticipation_length],
+                        noun_target[work_start: work_end + self.anticipation_length],
+                        segments_before_current,
+                    ])
+        else:
+            for session in self.sessions:
+                target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy'))
+                if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                    verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy'))
+                    noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy'))
+                else:
+                    verb_target = target
+                    noun_target = target
+                seed = np.random.randint(self.work_memory_length + self.anticipation_length) if self.training else 0
+                for work_start, work_end in zip(
+                    range(seed, target.shape[0], self.work_memory_length + self.anticipation_length),
+                    range(seed + self.work_memory_length, target.shape[0] - self.anticipation_length, self.work_memory_length + self.anticipation_length)):
+                    self.inputs.append([
+                        session, work_start, work_end,
+                        target[work_start: work_end + self.anticipation_length],
+                        verb_target[work_start: work_end + self.anticipation_length],
+                        noun_target[work_start: work_end + self.anticipation_length],
+                        None,
+                    ])
+                                                                                                             
+                                                                             
+                                      
+                                                                                  
+                    
+
+    def segment_sampler(self, start, end, num_samples):
+        indices = np.linspace(start, end, num_samples)
+        return np.sort(indices).astype(np.int32)
+
+    def uniform_sampler(self, start, end, num_samples, sample_rate):
+        if start < 0:
+            start = (end + 1) % sample_rate
+        indices = np.arange(start, end + 1)[::sample_rate]
+        padding = num_samples - indices.shape[0]
+        if padding > 0:
+            indices = np.concatenate((np.zeros(padding), indices))
+        return np.sort(indices).astype(np.int32)
+
+    def __getitem__(self, index):
+        (session, work_start, work_end, target,
+         verb_target, noun_target, segments_before) = self.inputs[index]
+        visual_inputs = np.load(
+            osp.join(self.data_root, self.visual_feature, session + '.npy'), mmap_mode='r')
+        motion_inputs = np.load(
+            osp.join(self.data_root, self.motion_feature, session + '.npy'), mmap_mode='r')
+        object_inputs = np.load(
+            osp.join(self.data_root, self.object_feature, session + '.npy'), mmap_mode='r')
+
+                    
+        target = np.concatenate((target[:self.work_memory_length:self.work_memory_sample_rate],
+                                 target[self.work_memory_length::self.anticipation_sample_rate]),
+                                 axis=0)
+        if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+            verb_target = np.concatenate((verb_target[:self.work_memory_length:self.work_memory_sample_rate],
+                                          verb_target[self.work_memory_length::self.anticipation_sample_rate]),
+                                          axis=0)
+            noun_target = np.concatenate((noun_target[:self.work_memory_length:self.work_memory_sample_rate],
+                                          noun_target[self.work_memory_length::self.anticipation_sample_rate]),
+                                          axis=0)
+
+                         
+        work_indices = np.arange(work_start, work_end).clip(0)
+        work_indices = work_indices[::self.work_memory_sample_rate]
+        work_visual_inputs = visual_inputs[work_indices]
+        work_motion_inputs = motion_inputs[work_indices]
+        work_object_inputs = object_inputs[work_indices]
+
+                         
+        if self.long_memory_num_samples > 0:
+                                                                                                 
+            long_start, long_end = work_start - self.long_memory_length, work_start - 1
+            if self.training:
+                long_indices = self.segment_sampler(
+                    long_start,
+                    long_end,
+                    self.long_memory_num_samples).clip(0)
+            else:
+                long_indices = self.uniform_sampler(
+                    long_start,
+                    long_end,
+                    self.long_memory_num_samples,
+                    self.long_memory_sample_rate).clip(0)
+            long_visual_inputs = visual_inputs[long_indices]
+            long_motion_inputs = motion_inputs[long_indices]
+            long_object_inputs = object_inputs[long_indices]
+                              
+            if self.training and self.clip_mixup_rate > 0:
+                assert segments_before is not None
+                segments_before = segments_before[segments_before['end_f'] / 30 * self.cfg.DATA.FPS > long_start]
+                if self.clip_mixup_sample == 'uniform':
+                    prob = None
+                elif self.clip_mixup_sample == 'by_length':
+                    prob = (segments_before['len_f']).tolist()
+                    prob = prob / np.sum(prob)
+                else:
+                    raise ValueError
+                num_clip = int(len(segments_before) * self.clip_mixup_rate)
+                segments_to_mixup = segments_before.sample(num_clip, replace=False,
+                                                           weights=prob)
+                for old_segment in segments_to_mixup.iterrows():
+                    old_start_tick = int(old_segment[1]['start_f'] / 30 * self.cfg.DATA.FPS)
+                    old_end_tick = int(np.ceil(old_segment[1]['end_f'] / 30 * self.cfg.DATA.FPS))
+                    old_action = old_segment[1]['action']
+                    old_vid = old_segment[1]['video']
+                    assert old_vid == session
+                    segment_same_class = self.segment_by_cls[old_action]
+                    segment_same_class = segment_same_class[segment_same_class['video'] != old_vid]
+                    valid_segments = segment_same_class[segment_same_class['len_f'] > old_segment[1]['len_f']]
+                    if len(valid_segments) == 0:
+                        continue
+                    else:
+                        sample_segment = valid_segments.sample(1)
+                        new_vid = sample_segment['video'].values[0]
+                        new_start_tick = int(sample_segment['start_f'] / 30 * self.cfg.DATA.FPS)
+                        new_end_tick = int(sample_segment['end_f'] / 30 * self.cfg.DATA.FPS)
+
+                        new_visual_inputs = np.load(
+                            osp.join(self.data_root, self.visual_feature, new_vid + '.npy'), mmap_mode='r')
+                        new_motion_inputs = np.load(
+                            osp.join(self.data_root, self.motion_feature, new_vid + '.npy'), mmap_mode='r')
+                        new_object_inputs = np.load(
+                            osp.join(self.data_root, self.object_feature, new_vid + '.npy'), mmap_mode='r')
+                        
+                        sel_indices = np.where((long_indices >= old_start_tick) & (long_indices <= old_end_tick))
+
+                        shift = np.random.randint(new_end_tick - new_start_tick - len(sel_indices) + 1)
+                        long_visual_inputs[sel_indices] = new_visual_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)]
+                        long_motion_inputs[sel_indices] = new_motion_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)]
+                        long_object_inputs[sel_indices] = new_object_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)]
+
+                                                      
+                                                                                                                                                                                      
+                                                      
+                                                                                                                                                                                      
+                                                      
+                                                                                                                                                                                      
+
+                                                                                                                   
+
+                                                                                                         
+                                                                                                             
+                                                                                                             
+                                                                                                             
+
+                                         
+            memory_key_padding_mask = np.zeros(long_indices.shape[0])
+            last_zero = bisect_right(long_indices, 0) - 1
+            if last_zero > 0:
+                memory_key_padding_mask[:last_zero] = float('-inf')
+        else:
+            long_visual_inputs = None
+            long_motion_inputs = None
+            long_object_inputs = None
+            memory_key_padding_mask = None
+
+                        
+        if (long_visual_inputs is not None and
+            long_motion_inputs is not None and
+            long_object_inputs is not None):
+            fusion_visual_inputs = np.concatenate((long_visual_inputs, work_visual_inputs))
+            fusion_motion_inputs = np.concatenate((long_motion_inputs, work_motion_inputs))
+            fusion_object_inputs = np.concatenate((long_object_inputs, work_object_inputs))
+        else:
+            fusion_visual_inputs = work_visual_inputs
+            fusion_motion_inputs = work_motion_inputs
+            fusion_object_inputs = work_object_inputs
+
+                           
+        fusion_visual_inputs = torch.as_tensor(fusion_visual_inputs.astype(np.float32))
+        fusion_motion_inputs = torch.as_tensor(fusion_motion_inputs.astype(np.float32))
+        fusion_object_inputs = torch.as_tensor(fusion_object_inputs.astype(np.float32))
+        target = torch.as_tensor(target.astype(np.float32))
+        if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                                                                    
+                                                             
+                                                           
+                                                                    
+                                                             
+                                                           
+                                                       
+                                                       
+                                                                   
+                                                                   
+                                                   
+                                                    
+                                                
+                                                          
+                                                   
+                                                    
+                                                
+                                                          
+            verb_target = torch.as_tensor(verb_target.astype(np.float32))
+            noun_target = torch.as_tensor(noun_target.astype(np.float32))
+            target = (target, verb_target, noun_target)
+
+        assert fusion_visual_inputs.ndim == fusion_motion_inputs.ndim
+        if fusion_visual_inputs.ndim == 4:
+                                    
+            fusion_visual_inputs = fusion_visual_inputs.permute(0, 2, 3, 1)
+            fusion_motion_inputs = fusion_motion_inputs.permute(0, 2, 3, 1)
+            fusion_visual_inputs = fusion_visual_inputs.mean((1, 2))
+            fusion_motion_inputs = fusion_motion_inputs.mean((1, 2))
+                                                                              
+                                                                                                          
+                                                                                                        
+        
+        if memory_key_padding_mask is not None:
+            memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32))
+            return (fusion_visual_inputs, fusion_motion_inputs,
+                    fusion_object_inputs, memory_key_padding_mask, target)
+        else:
+            return (fusion_visual_inputs, fusion_motion_inputs,
+                    fusion_object_inputs, fusion_object_inputs, target)
+
+    def __len__(self):
+        return len(self.inputs)
+
+
+@registry.register('LSTRBatchInferenceTHUMOS')
+@registry.register('LSTRBatchInferenceEGO4D')
+@registry.register('LSTRBatchInferenceEasyCom')
+@registry.register('LSTRBatchInferenceTVSeries')
+@registry.register('LSTRBatchInferenceEK55')
+@registry.register('LSTRBatchInferenceEK100')
+class LSTRBatchInferenceDataLayer(data.Dataset):
+
+    def __init__(self, cfg, phase='test'):
+        self.cfg = cfg
+        self.data_root = cfg.DATA.DATA_ROOT
+        self.visual_feature = cfg.INPUT.VISUAL_FEATURE
+        self.motion_feature = cfg.INPUT.MOTION_FEATURE
+        self.object_feature = cfg.INPUT.OBJECT_FEATURE
+        self.target_perframe = cfg.INPUT.TARGET_PERFRAME
+        self.sessions = getattr(cfg.DATA, phase.upper() + '_SESSION_SET')
+        self.long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH
+        self.long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE
+        self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES
+        self.work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH
+        self.work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE
+        self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES
+        self.anticipation_length = cfg.MODEL.LSTR.ANTICIPATION_LENGTH
+        self.anticipation_sample_rate = cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE
+        self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES
+
+        if phase != 'test':
+            warnings.warn('originally, phase must be `test` for batch inference, got {}'.format(phase))
+                                                                                    
+
+        self.inputs = []
+        for session in self.sessions:
+            target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy'))
+            if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy'))
+                noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy'))
+            else:
+                verb_target = target
+                noun_target = target
+            for work_start, work_end in zip(
+                range(0, target.shape[0] + 1),
+                range(self.work_memory_length, target.shape[0] + 1)):
+                self.inputs.append([
+                    session, work_start, work_end,
+                    target[work_start: work_end],
+                    verb_target[work_start: work_end],
+                    noun_target[work_start: work_end],
+                    target.shape[0]
+                ])
+
+    def uniform_sampler(self, start, end, num_samples, sample_rate):
+        if start < 0:
+            start = (end + 1) % sample_rate
+        indices = np.arange(start, end + 1)[::sample_rate]
+        padding = num_samples - indices.shape[0]
+        if padding > 0:
+            indices = np.concatenate((np.zeros(padding), indices))
+        return np.sort(indices).astype(np.int32)
+
+    def __getitem__(self, index):
+        session, work_start, work_end, target, verb_target, noun_target, num_frames = self.inputs[index]
+        visual_inputs = np.load(
+            osp.join(self.data_root, self.visual_feature, session + '.npy'), mmap_mode='r')
+        motion_inputs = np.load(
+            osp.join(self.data_root, self.motion_feature, session + '.npy'), mmap_mode='r')
+        object_inputs = np.load(
+            osp.join(self.data_root, self.object_feature, session + '.npy'), mmap_mode='r')
+
+                    
+        target = target[::self.work_memory_sample_rate]
+        verb_target = verb_target[::self.work_memory_sample_rate]
+        noun_target = noun_target[::self.work_memory_sample_rate]
+        
+                         
+        work_indices = np.arange(work_start, work_end).clip(0)
+        work_indices = work_indices[::self.work_memory_sample_rate]
+        work_visual_inputs = visual_inputs[work_indices]
+        work_motion_inputs = motion_inputs[work_indices]
+        work_object_inputs = object_inputs[work_indices]
+
+                         
+        if self.long_memory_num_samples > 0:
+                                                                                                 
+            long_start, long_end = work_start - self.long_memory_length, work_start - 1
+            long_indices = self.uniform_sampler(
+                long_start,
+                long_end,
+                self.long_memory_num_samples,
+                self.long_memory_sample_rate).clip(0)
+            long_visual_inputs = visual_inputs[long_indices]
+            long_motion_inputs = motion_inputs[long_indices]
+            long_object_inputs = object_inputs[long_indices]
+
+                                         
+            memory_key_padding_mask = np.zeros(long_indices.shape[0])
+            last_zero = bisect_right(long_indices, 0) - 1
+            if last_zero > 0:
+                memory_key_padding_mask[:last_zero] = float('-inf')
+        else:
+            long_visual_inputs = None
+            long_motion_inputs = None
+            long_object_inputs = None
+            memory_key_padding_mask = None
+
+                        
+        if long_visual_inputs is not None and long_motion_inputs is not None:
+            fusion_visual_inputs = np.concatenate((long_visual_inputs, work_visual_inputs))
+            fusion_motion_inputs = np.concatenate((long_motion_inputs, work_motion_inputs))
+            fusion_object_inputs = np.concatenate((long_object_inputs, work_object_inputs))
+        else:
+            fusion_visual_inputs = work_visual_inputs
+            fusion_motion_inputs = work_motion_inputs
+            fusion_object_inputs = work_object_inputs
+
+                           
+        fusion_visual_inputs = torch.as_tensor(fusion_visual_inputs.astype(np.float32))
+        fusion_motion_inputs = torch.as_tensor(fusion_motion_inputs.astype(np.float32))
+        fusion_object_inputs = torch.as_tensor(fusion_object_inputs.astype(np.float32))
+        target = torch.as_tensor(target.astype(np.float32))
+
+        if self.cfg.MODEL.LSTR.V_N_CLASSIFIER:
+            verb_target = torch.as_tensor(verb_target.astype(np.float32))
+            noun_target = torch.as_tensor(noun_target.astype(np.float32))
+            target = (target, verb_target, noun_target)
+
+        if memory_key_padding_mask is not None:
+            memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32))
+            return (fusion_visual_inputs, fusion_motion_inputs,
+                    fusion_object_inputs, memory_key_padding_mask, target,
+                    session, work_indices, num_frames)
+        else:
+            return (fusion_visual_inputs, fusion_motion_inputs,
+                    fusion_object_inputs, target,
+                    session, work_indices, num_frames)
+
+    def __len__(self):
+        return len(self.inputs)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/__init__.py

@@ -0,0 +1,6 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .engines import do_inference, do_train
+from .lstr.lstr_inference import *
+from .lstr.lstr_trainer import *

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/perframe_det_batch_inference.py

@@ -0,0 +1,673 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import os.path as osp
+import pickle as pkl
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+from rekognition_online_action_detection.datasets import build_dataset
+from rekognition_online_action_detection.evaluation import compute_result
+from tqdm import tqdm
+
+
+def batch_inference(cfg, model, device, logger, data_loader):
+    
+    # Collect scores and targets
+    pred_scores = {}
+    pred_scores_verb = {}
+    pred_scores_noun = {}
+    gt_targets = {}
+    
+    start_time = time.time()
+    with torch.no_grad():
+        pbar = tqdm(data_loader, desc='BatchInference')
+        for batch_idx, data in enumerate(pbar, start=1):
+            target = data[-4]
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                target, verb_target, noun_target = target
+
+            score = model(*[x.to(device) for x in data[:-4]])
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                score, score_verb, score_noun = score
+                score = score.cpu().numpy()
+                score_verb = score_verb.cpu().numpy()
+                score_noun = score_noun.cpu().numpy()
+                cfg.DATA.NUM_VERBS = 126 if cfg.DATA.DATA_NAME == 'EK55' else 98
+                cfg.DATA.NUM_NOUNS = 353 if cfg.DATA.DATA_NAME == 'EK55' else 301
+            else:
+                score = (score.cpu().numpy()
+                         if cfg.DATA.DATA_NAME.startswith('EK')
+                         else score.softmax(dim=-1).cpu().numpy())
+            for bs, (session, query_indices, num_frames) in enumerate(zip(*data[-3:])):
+                # data[-3] is session, data[-2] is work_indices, data[-1] is num_frames
+                if session not in pred_scores:
+                    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                        pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+                        if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                            pred_scores_verb[session] = np.zeros((num_frames, cfg.DATA.NUM_VERBS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+                            pred_scores_noun[session] = np.zeros((num_frames, cfg.DATA.NUM_NOUNS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+                    else:
+                        pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES))
+                if session not in gt_targets:
+                    gt_targets[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES))
+
+                # DEFAULT VALUES in config
+                # cfg.MODEL.LSTR.WORK_MEMORY_LENGTH: 40
+                # cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE: 1
+                # cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES: 10
+                # query_indices.shape: torch.Size([cfg.MODEL.LSTR.WORK_MEMORY_LENGTH])
+                # query_indices at first: [0, 1, 2, 3, ..., 39]
+                # At first, go into this if statement
+                if query_indices[0] in torch.arange(0, cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE):
+                    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                            # at first, full_indices indicate the range:
+                            # work_memory + t_a (anticiptaion timestep, from 0 to 9)
+                            full_indices = torch.cat((query_indices,
+                                                     torch.arange(query_indices[-1] + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE,
+                                                                  query_indices[-1] + (t_a + 1) + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE)),
+                                                     dim=0)
+                            
+                            # then assgin the currnet logits + anticipated logits to the pred_scores
+                            pred_scores[session][full_indices, :, t_a] = score[bs][:full_indices.shape[0]]
+                            
+                            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                                pred_scores_verb[session][full_indices, :, t_a] = score_verb[bs][:full_indices.shape[0]]    
+                                pred_scores_noun[session][full_indices, :, t_a] = score_noun[bs][:full_indices.shape[0]]    
+                                
+                        gt_targets[session][query_indices] = target[bs][:query_indices.shape[0]]
+                    else:
+                        # no anticipation
+                        pred_scores[session][query_indices] = score[bs]
+                        gt_targets[session][query_indices] = target[bs]
+                else:
+                    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                            if query_indices[-1] + t_a + 1 < num_frames:
+                                pred_scores[session][query_indices[-1] + t_a + 1, :, t_a] = score[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]
+                                if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                                    pred_scores_verb[session][query_indices[-1] + t_a + 1, :, t_a] = score_verb[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]    
+                                    pred_scores_noun[session][query_indices[-1] + t_a + 1, :, t_a] = score_noun[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]
+                        gt_targets[session][query_indices[-1]] = target[bs][-1]
+                    else:
+                        pred_scores[session][query_indices[-1]] = score[bs][-1]
+                        gt_targets[session][query_indices[-1]] = target[bs][-1]
+    time_taken = time.time() - start_time
+    # Save scores and targets
+    # pkl.dump({
+    #     'cfg': cfg,
+    #     'perframe_pred_scores': pred_scores,
+    #     'perframe_gt_targets': gt_targets,
+    # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.pkl', 'wb'))
+    num_frames = np.concatenate(list(gt_targets.values()), axis=0).shape[0]
+    logger.info(f'\nProcessed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)')
+    logger.info('\n=============')
+    # Compute results
+    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+        maps_list = []
+        per_class_list = {}
+
+        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+            result = compute_result['perframe'](
+                cfg,
+                np.concatenate(list(gt_targets.values()), axis=0),
+                np.concatenate(list(pred_scores.values()), axis=0)[:, :, t_a],
+            )
+            
+            logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format(
+                (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+                cfg.DATA.METRICS, result['mean_AP']
+            ))
+            maps_list.append(result['mean_AP'])
+            # per_class_AP = result['per_class_AP']
+            
+            anticipation_time=(t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE
+            
+            for class_name, ap in result['per_class_AP'].items():
+                if class_name not in per_class_list:
+                    per_class_list[class_name] = []
+                per_class_list[class_name].append(ap)
+                logger.info(f'Action anticipation ({anticipation_time:.2f}s) perframe AP of {class_name}: {ap:.5f}')
+            
+
+        logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format(
+            cfg.DATA.METRICS, np.mean(maps_list)
+        ))
+        
+        for class_name, ap_list in per_class_list.items():
+            logger.info(f'Action anticipation (mean) perframe AP of {class_name}: {np.mean(ap_list):.5f}')
+        
+        # EDIT: save predictions!
+        for session in pred_scores.keys():
+            anticipation_score = pred_scores[session]
+            anticipation_dir = os.path.join(cfg.OUTPUT_DIR, 'pred')
+            if not os.path.exists(anticipation_dir):
+                os.makedirs(anticipation_dir)
+            
+            np.save(osp.join(anticipation_dir, session+'.npy'), anticipation_score)
+        
+    else:
+        result = compute_result['perframe'](
+            cfg,
+            np.concatenate(list(gt_targets.values()), axis=0),
+            np.concatenate(list(pred_scores.values()), axis=0),
+        )
+        logger.info('Action detection perframe m{}: {:.5f}'.format(
+            cfg.DATA.METRICS, result['mean_AP']
+        ))
+
+    # segment-level evaluation on EK55/EK100
+    if cfg.DATA.DATA_NAME.startswith('EK'):
+        try:
+            from external.rulstm.RULSTM.utils import (
+                get_marginal_indexes, marginalize, softmax,
+                topk_accuracy_multiple_timesteps,
+                topk_recall_multiple_timesteps, tta)
+        except:
+            raise ModuleNotFoundError
+        if cfg.DATA.DATA_NAME == 'EK55':
+            path_to_data = 'external/rulstm/RULSTM/data/ek55/'
+        elif cfg.DATA.DATA_NAME == 'EK100':
+            path_to_data = 'external/rulstm/RULSTM/data/ek100/'
+        segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'),
+                                   names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+                                   skipinitialspace=True)
+        predictions, labels, ids = [], [], []
+        predictions_verb, predictions_noun = [], []
+        discarded_labels, discarded_ids = [], []
+        for segment in segment_list.iterrows():
+            start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS))
+            vid = segment[1]['video'].strip()
+            preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+                preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+            else:
+                preds_v = None
+                preds_n = None
+            if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH:
+                # pt = []
+                # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_LENGTH):
+                #     pt.append(pred_scores[vid][start_step - t_a, 1:, t_a][None, None, :])
+                # preds = np.concatenate(pt, axis=1)
+                predictions.append(preds)
+                if preds_v is not None:
+                    predictions_verb.append(preds_v)
+                if preds_n is not None:
+                    predictions_noun.append(preds_n)
+                labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+                ids.append(segment[1]['id'])
+            else:
+                discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+                discarded_ids.append(segment[1]['id'])
+        
+        actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id')
+
+        action_scores = np.concatenate(predictions)
+        if len(predictions_verb) > 0:
+            verb_scores = np.concatenate(predictions_verb)
+        else:
+            verb_scores = None
+        if len(predictions_noun) > 0:
+            noun_scores = np.concatenate(predictions_noun)
+        else:
+            noun_scores = None
+        labels = np.array(labels)
+        ids = np.array(ids)
+
+        vi = get_marginal_indexes(actions, 'verb')
+        ni = get_marginal_indexes(actions, 'noun')
+
+        pkl.dump({
+            'cfg': cfg,
+            'action_scores': action_scores,
+            'verb_scores': verb_scores,
+            'noun_scores': noun_scores,
+        }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.ek55_style.pkl', 'wb'))
+
+        action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1]))
+        # action_probs = action_probs / action_probs.sum(-1, keepdims=True)
+        print(action_probs.shape)
+        if verb_scores is None:
+            verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0],
+                                                                action_scores.shape[1],
+                                                                -1)
+        else:
+            verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1]))
+            verb_scores = verb_scores.reshape(action_scores.shape[0],
+                                              action_scores.shape[1],
+                                              -1)
+        if noun_scores is None:
+            noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0],
+                                                                action_scores.shape[1],
+                                                                -1)
+        else:
+            noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1]))
+            noun_scores = noun_scores.reshape(action_scores.shape[0],
+                                              action_scores.shape[1],
+                                              -1)
+
+        include_discarded = (cfg.DATA.DATA_NAME == 'EK100')
+        if include_discarded:
+            dlab = np.array(discarded_labels)
+            dislab = np.array(discarded_ids)
+            ids = np.concatenate([ids, dislab])
+            num_disc = len(dlab)
+            labels = np.concatenate([labels, dlab])
+            verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:]))))
+            noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:]))))
+            action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:]))))
+
+        verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2]
+        verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels)
+        noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels)
+        action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels)
+
+        if cfg.DATA.DATA_NAME == 'EK55':
+            many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values
+            many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values
+            actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'))
+            a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values)
+                    for a in actions.iterrows()}
+            many_shot_actions = [a for a, (v, n) in a_to_vn.items()
+                                if v in many_shot_verbs or n in many_shot_nouns]
+            
+            verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels,
+                                                        k=5, classes=many_shot_verbs)
+            noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels,
+                                                        k=5, classes=many_shot_nouns)
+            action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels,
+                                                            k=5, classes=many_shot_actions)
+            
+            all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies,
+                                            verb_recalls, noun_recalls, action_recalls])
+            all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]]
+            indices = [
+                ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'),
+                ('Verb', 'Mean Top-5 Recall'),
+                ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'),
+                ('Noun', 'Mean Top-5 Recall'),
+                ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'),
+                ('Action', 'Mean Top-5 Recall'),
+            ]
+        else:
+            overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5)
+            overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5)
+            overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5)
+
+            unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True)
+            tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True)
+            tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True)
+            tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True)
+
+            unseen_bool_idx = pd.Series(ids).isin(unseen).values
+            tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values
+            tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values
+            tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values
+
+            tail_verb_recalls = topk_recall_multiple_timesteps(
+                verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5)
+            tail_noun_recalls = topk_recall_multiple_timesteps(
+                noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5)
+            tail_action_recalls = topk_recall_multiple_timesteps(
+                action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5)
+            
+            unseen_verb_recalls = topk_recall_multiple_timesteps(
+                verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5)
+            unseen_noun_recalls = topk_recall_multiple_timesteps(
+                noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5)
+            unseen_action_recalls = topk_recall_multiple_timesteps(
+                action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5)
+
+            all_accuracies = np.concatenate(
+                [overall_verb_recalls, overall_noun_recalls, overall_action_recalls,
+                 unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls,
+                 tail_verb_recalls, tail_noun_recalls, tail_action_recalls]
+            ) #9 x 8
+
+            indices = [
+                ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'),
+                ('Overall Mean Top-5 Recall', 'Action'),
+                ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'),
+                ('Unseen Mean Top-5 Recall', 'Action'),
+                ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'),
+                ('Tail Mean Top-5 Recall', 'Action'),
+            ]
+
+        cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH,
+                        1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+                        cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str)
+        scores = pd.DataFrame(all_accuracies * 100, columns=cc,
+                            index=pd.MultiIndex.from_tuples(indices))
+        logger.info(scores)
+
+        tta_verb = tta(verb_scores, verb_labels)
+        tta_noun = tta(noun_scores, noun_labels)
+        tta_action = tta(action_scores, action_labels)
+        logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} '
+                    f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}')
+
+
+def do_perframe_det_batch_inference(cfg, model, device, logger):
+    # Setup model to test mode
+    model.eval()
+    cfg.MODEL.CRITERIONS = [['MCE', {}]]
+    
+    data_loader = torch.utils.data.DataLoader(
+        dataset=build_dataset(cfg, phase='test', tag='BatchInference'),
+        batch_size=cfg.DATA_LOADER.BATCH_SIZE,
+        num_workers=cfg.DATA_LOADER.NUM_WORKERS,
+        pin_memory=cfg.DATA_LOADER.PIN_MEMORY,
+    )
+
+    logger.info(f"phase: test) Inference on test_dataloader")
+    start = time.time()
+    batch_inference(cfg, model, device, logger, data_loader)
+    end = time.time() - start
+    
+    logger.info(f"Elapsed time for function batch_inference: {end:.2f}s")
+    
+    # data_loader = torch.utils.data.DataLoader(
+    #     dataset=build_dataset(cfg, phase='train', tag='BatchInference'),
+    #     batch_size=cfg.DATA_LOADER.BATCH_SIZE,
+    #     num_workers=cfg.DATA_LOADER.NUM_WORKERS,
+    #     pin_memory=cfg.DATA_LOADER.PIN_MEMORY,
+    # )
+    
+    # logger.info(f"***** phase: train")
+    # batch_inference(cfg, model, device, logger, data_loader)
+    
+    # # Collect scores and targets
+    # pred_scores = {}
+    # pred_scores_verb = {}
+    # pred_scores_noun = {}
+    # gt_targets = {}
+
+    # with torch.no_grad():
+    #     pbar = tqdm(data_loader, desc='BatchInference')
+    #     for batch_idx, data in enumerate(pbar, start=1):
+    #         target = data[-4]
+    #         if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #             target, verb_target, noun_target = target
+
+    #         score = model(*[x.to(device) for x in data[:-4]])
+    #         if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #             score, score_verb, score_noun = score
+    #             score = score.cpu().numpy()
+    #             score_verb = score_verb.cpu().numpy()
+    #             score_noun = score_noun.cpu().numpy()
+    #             cfg.DATA.NUM_VERBS = 126 if cfg.DATA.DATA_NAME == 'EK55' else 98
+    #             cfg.DATA.NUM_NOUNS = 353 if cfg.DATA.DATA_NAME == 'EK55' else 301
+    #         else:
+    #             score = (score.cpu().numpy()
+    #                      if cfg.DATA.DATA_NAME.startswith('EK')
+    #                      else score.softmax(dim=-1).cpu().numpy())
+    #         for bs, (session, query_indices, num_frames) in enumerate(zip(*data[-3:])):
+    #             if session not in pred_scores:
+    #                 if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+    #                     pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+    #                     if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #                         pred_scores_verb[session] = np.zeros((num_frames, cfg.DATA.NUM_VERBS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+    #                         pred_scores_noun[session] = np.zeros((num_frames, cfg.DATA.NUM_NOUNS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES))
+    #                 else:
+    #                     pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES))
+    #             if session not in gt_targets:
+    #                 gt_targets[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES))
+
+    #             if query_indices[0] in torch.arange(0, cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE):
+    #                 if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+    #                     for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+    #                         full_indices = torch.cat((query_indices,
+    #                                                  torch.arange(query_indices[-1] + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE,
+    #                                                               query_indices[-1] + t_a + 1 + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE)),
+    #                                                  dim=0)
+    #                         pred_scores[session][full_indices, :, t_a] = score[bs][:full_indices.shape[0]]
+    #                         if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #                             pred_scores_verb[session][full_indices, :, t_a] = score_verb[bs][:full_indices.shape[0]]    
+    #                             pred_scores_noun[session][full_indices, :, t_a] = score_noun[bs][:full_indices.shape[0]]    
+    #                     gt_targets[session][query_indices] = target[bs][:query_indices.shape[0]]
+    #                 else:
+    #                     pred_scores[session][query_indices] = score[bs]
+    #                     gt_targets[session][query_indices] = target[bs]
+    #             else:
+    #                 if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+    #                     for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+    #                         if query_indices[-1] + t_a + 1 < num_frames:
+    #                             pred_scores[session][query_indices[-1] + t_a + 1, :, t_a] = score[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]
+    #                             if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #                                 pred_scores_verb[session][query_indices[-1] + t_a + 1, :, t_a] = score_verb[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]    
+    #                                 pred_scores_noun[session][query_indices[-1] + t_a + 1, :, t_a] = score_noun[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]
+    #                     gt_targets[session][query_indices[-1]] = target[bs][-1]
+    #                 else:
+    #                     pred_scores[session][query_indices[-1]] = score[bs][-1]
+    #                     gt_targets[session][query_indices[-1]] = target[bs][-1]
+
+    # # Save scores and targets
+    # # pkl.dump({
+    # #     'cfg': cfg,
+    # #     'perframe_pred_scores': pred_scores,
+    # #     'perframe_gt_targets': gt_targets,
+    # # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.pkl', 'wb'))
+
+    # # Compute results
+    # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+    #     maps_list = []
+    #     for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+    #         result = compute_result['perframe'](
+    #             cfg,
+    #             np.concatenate(list(gt_targets.values()), axis=0),
+    #             np.concatenate(list(pred_scores.values()), axis=0)[:, :, t_a],
+    #         )
+    #         logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format(
+    #             (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+    #             cfg.DATA.METRICS, result['mean_AP']
+    #         ))
+    #         maps_list.append(result['mean_AP'])
+    #     logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format(
+    #         cfg.DATA.METRICS, np.mean(maps_list)
+    #     ))
+    # else:
+    #     result = compute_result['perframe'](
+    #         cfg,
+    #         np.concatenate(list(gt_targets.values()), axis=0),
+    #         np.concatenate(list(pred_scores.values()), axis=0),
+    #     )
+    #     logger.info('Action detection perframe m{}: {:.5f}'.format(
+    #         cfg.DATA.METRICS, result['mean_AP']
+    #     ))
+
+    # # segment-level evaluation on EK55/EK100
+    # if cfg.DATA.DATA_NAME.startswith('EK'):
+    #     try:
+    #         from external.rulstm.RULSTM.utils import (get_marginal_indexes, marginalize, softmax,
+    #                                                   topk_accuracy_multiple_timesteps,
+    #                                                   topk_recall_multiple_timesteps,
+    #                                                   tta)
+    #     except:
+    #         raise ModuleNotFoundError
+    #     if cfg.DATA.DATA_NAME == 'EK55':
+    #         path_to_data = 'external/rulstm/RULSTM/data/ek55/'
+    #     elif cfg.DATA.DATA_NAME == 'EK100':
+    #         path_to_data = 'external/rulstm/RULSTM/data/ek100/'
+    #     segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'),
+    #                                names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+    #                                skipinitialspace=True)
+    #     predictions, labels, ids = [], [], []
+    #     predictions_verb, predictions_noun = [], []
+    #     discarded_labels, discarded_ids = [], []
+    #     for segment in segment_list.iterrows():
+    #         start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS))
+    #         vid = segment[1]['video'].strip()
+    #         preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+    #         if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+    #             preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+    #             preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+    #         else:
+    #             preds_v = None
+    #             preds_n = None
+    #         if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH:
+    #             # pt = []
+    #             # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_LENGTH):
+    #             #     pt.append(pred_scores[vid][start_step - t_a, 1:, t_a][None, None, :])
+    #             # preds = np.concatenate(pt, axis=1)
+    #             predictions.append(preds)
+    #             if preds_v is not None:
+    #                 predictions_verb.append(preds_v)
+    #             if preds_n is not None:
+    #                 predictions_noun.append(preds_n)
+    #             labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+    #             ids.append(segment[1]['id'])
+    #         else:
+    #             discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+    #             discarded_ids.append(segment[1]['id'])
+        
+    #     actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id')
+
+    #     action_scores = np.concatenate(predictions)
+    #     if len(predictions_verb) > 0:
+    #         verb_scores = np.concatenate(predictions_verb)
+    #     else:
+    #         verb_scores = None
+    #     if len(predictions_noun) > 0:
+    #         noun_scores = np.concatenate(predictions_noun)
+    #     else:
+    #         noun_scores = None
+    #     labels = np.array(labels)
+    #     ids = np.array(ids)
+
+    #     vi = get_marginal_indexes(actions, 'verb')
+    #     ni = get_marginal_indexes(actions, 'noun')
+
+    #     pkl.dump({
+    #         'cfg': cfg,
+    #         'action_scores': action_scores,
+    #         'verb_scores': verb_scores,
+    #         'noun_scores': noun_scores,
+    #     }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.ek55_style.pkl', 'wb'))
+
+    #     action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1]))
+    #     # action_probs = action_probs / action_probs.sum(-1, keepdims=True)
+    #     print(action_probs.shape)
+    #     if verb_scores is None:
+    #         verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0],
+    #                                                             action_scores.shape[1],
+    #                                                             -1)
+    #     else:
+    #         verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1]))
+    #         verb_scores = verb_scores.reshape(action_scores.shape[0],
+    #                                           action_scores.shape[1],
+    #                                           -1)
+    #     if noun_scores is None:
+    #         noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0],
+    #                                                             action_scores.shape[1],
+    #                                                             -1)
+    #     else:
+    #         noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1]))
+    #         noun_scores = noun_scores.reshape(action_scores.shape[0],
+    #                                           action_scores.shape[1],
+    #                                           -1)
+
+    #     include_discarded = (cfg.DATA.DATA_NAME == 'EK100')
+    #     if include_discarded:
+    #         dlab = np.array(discarded_labels)
+    #         dislab = np.array(discarded_ids)
+    #         ids = np.concatenate([ids, dislab])
+    #         num_disc = len(dlab)
+    #         labels = np.concatenate([labels, dlab])
+    #         verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:]))))
+    #         noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:]))))
+    #         action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:]))))
+
+    #     verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2]
+    #     verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels)
+    #     noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels)
+    #     action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels)
+
+    #     if cfg.DATA.DATA_NAME == 'EK55':
+    #         many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values
+    #         many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values
+    #         actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'))
+    #         a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values)
+    #                 for a in actions.iterrows()}
+    #         many_shot_actions = [a for a, (v, n) in a_to_vn.items()
+    #                             if v in many_shot_verbs or n in many_shot_nouns]
+            
+    #         verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels,
+    #                                                     k=5, classes=many_shot_verbs)
+    #         noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels,
+    #                                                     k=5, classes=many_shot_nouns)
+    #         action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels,
+    #                                                         k=5, classes=many_shot_actions)
+            
+    #         all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies,
+    #                                         verb_recalls, noun_recalls, action_recalls])
+    #         all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]]
+    #         indices = [
+    #             ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'),
+    #             ('Verb', 'Mean Top-5 Recall'),
+    #             ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'),
+    #             ('Noun', 'Mean Top-5 Recall'),
+    #             ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'),
+    #             ('Action', 'Mean Top-5 Recall'),
+    #         ]
+    #     else:
+    #         overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5)
+    #         overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5)
+    #         overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5)
+
+    #         unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True)
+    #         tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True)
+    #         tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True)
+    #         tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True)
+
+    #         unseen_bool_idx = pd.Series(ids).isin(unseen).values
+    #         tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values
+    #         tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values
+    #         tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values
+
+    #         tail_verb_recalls = topk_recall_multiple_timesteps(
+    #             verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5)
+    #         tail_noun_recalls = topk_recall_multiple_timesteps(
+    #             noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5)
+    #         tail_action_recalls = topk_recall_multiple_timesteps(
+    #             action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5)
+            
+    #         unseen_verb_recalls = topk_recall_multiple_timesteps(
+    #             verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5)
+    #         unseen_noun_recalls = topk_recall_multiple_timesteps(
+    #             noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5)
+    #         unseen_action_recalls = topk_recall_multiple_timesteps(
+    #             action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5)
+
+    #         all_accuracies = np.concatenate(
+    #             [overall_verb_recalls, overall_noun_recalls, overall_action_recalls,
+    #              unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls,
+    #              tail_verb_recalls, tail_noun_recalls, tail_action_recalls]
+    #         ) #9 x 8
+
+    #         indices = [
+    #             ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'),
+    #             ('Overall Mean Top-5 Recall', 'Action'),
+    #             ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'),
+    #             ('Unseen Mean Top-5 Recall', 'Action'),
+    #             ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'),
+    #             ('Tail Mean Top-5 Recall', 'Action'),
+    #         ]
+
+    #     cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH,
+    #                     1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+    #                     cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str)
+    #     scores = pd.DataFrame(all_accuracies * 100, columns=cc,
+    #                         index=pd.MultiIndex.from_tuples(indices))
+    #     logger.info(scores)
+
+    #     tta_verb = tta(verb_scores, verb_labels)
+    #     tta_noun = tta(noun_scores, noun_labels)
+    #     tta_action = tta(action_scores, action_labels)
+    #     logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} '
+    #                 f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}')
+    #                 f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}')

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer.py

@@ -0,0 +1,245 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import os.path as osp
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import wandb
+from rekognition_online_action_detection.evaluation import compute_result
+from tqdm import tqdm
+
+
+def log_results(phase, cfg, compute_result, gt_targets, pred_scores, losses, data_loader, epoch=None, best_mAP=None, best_epoch=None):
+    result = compute_result['perframe'](cfg, gt_targets, pred_scores)
+    per_class_AP = result['mean_AP']
+    
+    per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items())
+    
+    log = []
+    if epoch is not None:
+        log.append(f'\nEpoch {epoch:2}')
+    
+    log.append(f'{phase} det_loss: {losses["det"] / len(data_loader.dataset):.5f} det_mAP: {result["mean_AP"]:.5f} {per_class_AP_log}')
+    
+    if cfg.USE_WANDB:
+        wandb.log({
+            f"{phase.capitalize()} mAP": result['mean_AP'],
+            f"{phase.capitalize()} per_class_AP": per_class_AP,
+            f"{phase.capitalize()} Loss": losses['det'] / len(data_loader.dataset)
+        })
+    
+    if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+        for loss_type in ['verb', 'noun']:
+            if loss_type in losses:
+                log.append(f'{phase} {loss_type}_loss: {losses[loss_type] / len(data_loader.dataset):.5f}')
+    
+    return log, result['mean_AP']
+
+def do_perframe_det_train(cfg,
+                          data_loaders,
+                          model,
+                          criterion,
+                          optimizer,
+                          scheduler,
+                          device,
+                          checkpointer,
+                          logger):
+    # Setup model on multiple GPUs
+    if torch.cuda.device_count() > 1:
+        model = nn.DataParallel(model)
+
+    best_mAP, best_epoch = 0, 0
+    best_per_class_AP_log = None
+
+    for epoch in range(cfg.SOLVER.START_EPOCH, cfg.SOLVER.START_EPOCH + cfg.SOLVER.NUM_EPOCHS):
+        # Reset
+        det_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        pred_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        verb_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        noun_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        det_pred_scores = []
+        det_gt_targets = []
+        aa_pred_scores = []
+        aa_gt_targets = []
+        
+        train_det_pred_scores = []
+        train_det_gt_targets = []
+        train_aa_pred_scores = []
+        train_aa_gt_targets = []
+
+        start = time.time()
+        
+        for phase in cfg.SOLVER.PHASES:
+            training = phase == 'train'
+            model.train(training)
+
+            with torch.set_grad_enabled(training):
+                pbar = tqdm(data_loaders[phase],
+                            desc='{}ing epoch {}'.format(phase.capitalize(), epoch))
+                for batch_idx, data in enumerate(pbar, start=1):
+                    batch_size = data[0].shape[0]
+                    
+                    # batch_size 16
+                    # data[0]: visual_input, [16, 680, 2048]
+                    # [BATCH_SIZE, 640 + 40 (LONG_MEMORY_NUM_SAMPLES + WORK_MEMORY_NUM_SAMPLES), FEATURE_SIZE]
+                    # this if statement is for the case of EK100 dataset
+                    if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                        det_target, verb_target, noun_target = data[-1]
+                        det_target = det_target.to(device)
+                        verb_target = verb_target.to(device)
+                        noun_target = noun_target.to(device)
+                    else:
+                        det_target = data[-1].to(device)
+
+                    loss_names = list(zip(*cfg.MODEL.CRITERIONS))[0]
+                    if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]:
+                        det_score, feat_out, feat_ori = model(*[x.to(device) for x in data[:-1]])
+                    else:
+                        det_score = model(*[x.to(device) for x in data[:-1]])
+
+                    if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                        det_score, verb_score, noun_score = det_score
+                    else:
+                        verb_score, noun_score = None, None
+                        
+                    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                        aa_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                        # [16, 10, 3]
+                        aa_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                        
+                    if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY:
+                        det_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                        det_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                    
+                    # det_score: [16, 50, 3], aa_score: [16, 10, 3]
+                    # det_score can be different when aa_score 
+                                       
+                    det_score = det_score.reshape(-1, cfg.DATA.NUM_CLASSES)
+                    det_target = det_target.reshape(-1, cfg.DATA.NUM_CLASSES)
+                    
+                    assert len(det_score.shape) == 2                    
+                    det_loss = criterion[loss_names[0]](det_score, det_target)
+                    det_losses[phase] += det_loss.item() * batch_size
+
+                    if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]:
+                        pred_loss = criterion['PRED_FUTURE'](feat_out, feat_ori)
+                        pred_losses[phase] += pred_loss.item() * batch_size
+
+                    if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                        if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY:
+                            verb_score = verb_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                            verb_target = verb_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                            noun_score = noun_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                            noun_target = noun_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                        verb_score = verb_score.reshape(-1, verb_score.shape[-1])
+                        noun_score = noun_score.reshape(-1, noun_score.shape[-1])
+                        verb_target = verb_target.reshape(-1, verb_target.shape[-1])
+                        noun_target = noun_target.reshape(-1, noun_target.shape[-1])
+                        verb_loss = criterion['MCE'](verb_score, verb_target)
+                        noun_loss = criterion['MCE'](noun_score, noun_target)
+                        verb_losses[phase] += verb_loss.item() * batch_size
+                        noun_losses[phase] += noun_loss.item() * batch_size
+
+                    # Output log for current batch
+                    pbar.set_postfix({
+                        'lr': '{:.7f}'.format(scheduler.get_last_lr()[0]),
+                        'det_loss': '{:.5f}'.format(det_loss.item()),
+                    })
+
+                    if training:
+                        optimizer.zero_grad()
+                        loss = det_loss
+                        if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]:
+                            loss += pred_loss
+                        if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                            loss += noun_loss + verb_loss
+                        if loss.item() != 0:
+                            loss.backward()
+                            optimizer.step()
+                            scheduler.step()
+                        
+                        if cfg.USE_WANDB:
+                            wandb.log({"Learning Rate": scheduler.get_last_lr()[0]})
+                            wandb.log({"Train Loss": det_loss.item()})
+                        
+                        det_score = det_score.softmax(dim=1).cpu().tolist()
+                        det_target = det_target.cpu().tolist()
+                        train_det_pred_scores.extend(det_score)
+                        train_det_gt_targets.extend(det_target)
+                        
+                        if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                            aa_score = aa_score.softmax(dim=1).cpu().tolist()
+                            aa_target = aa_target.cpu().tolist()
+                            train_aa_pred_scores.extend(aa_score)
+                            train_aa_gt_targets.extend(aa_target)
+                            
+                    else:
+                        # Prepare for evaluation
+                        det_score = det_score.softmax(dim=1).cpu().tolist()
+                        det_target = det_target.cpu().tolist()
+                        det_pred_scores.extend(det_score)
+                        det_gt_targets.extend(det_target)
+                        
+                        if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                            aa_score = aa_score.softmax(dim=1).cpu().tolist()
+                            aa_target = aa_target.cpu().tolist()
+                            aa_pred_scores.extend(aa_score)
+                            aa_gt_targets.extend(aa_target)
+                        
+        end = time.time()
+
+        # After the phase loop, compute and log results
+        for phase in cfg.SOLVER.PHASES:
+            if phase == 'train':
+                gt_targets = train_det_gt_targets
+                pred_scores = train_det_pred_scores
+            else:  # 'test'
+                gt_targets = det_gt_targets
+                pred_scores = det_pred_scores
+            
+            result = compute_result['perframe'](cfg, gt_targets, pred_scores)
+            per_class_AP = result['per_class_AP']
+            
+            log = []
+            log.append(f'\nEpoch {epoch:2}')
+            log.append(f'{phase} det_loss: {det_losses[phase] / len(data_loaders[phase].dataset):.5f} det_mAP: {result["mean_AP"]:.5f}')
+            
+            per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items())
+            log.append(f'Per-class AP: {per_class_AP_log}')
+            
+            if cfg.USE_WANDB:
+                wandb.log({
+                    f"{phase.capitalize()} mAP": result['mean_AP'],
+                    f"{phase.capitalize()} per_class_AP": per_class_AP,
+                    f"{phase.capitalize()} Loss": det_losses[phase] / len(data_loaders[phase].dataset)
+                })
+            
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                log.append(f'{phase} verb_loss: {verb_losses[phase] / len(data_loaders[phase].dataset):.5f}')
+                log.append(f'{phase} noun_loss: {noun_losses[phase] / len(data_loaders[phase].dataset):.5f}')
+            
+            if phase == 'test':
+                if result['mean_AP'] > best_mAP:
+                    best_mAP = result['mean_AP']
+                    best_epoch = epoch
+                    best_per_class_AP_log = per_class_AP_log
+                    checkpointer.save_best(epoch, model, optimizer)
+                
+                log.append(f'\nbest epoch: {best_epoch:2} best det_mAP: {best_mAP:.5f} best per-class AP: {best_per_class_AP_log}\n')
+            
+            logger.info(' | '.join(log))
+
+        log.append(f'running time: {end - start:.2f} sec')
+        logger.info(' | '.join(log))
+
+        # Save checkpoint for model and optimizer
+        if cfg.SOLVER.SAVE_EVERY_EPOCHS > 0 and epoch % cfg.SOLVER.SAVE_EVERY_EPOCHS == 0:
+            checkpointer.save(epoch, model, optimizer)
+
+        # Shuffle dataset for next epoch
+        data_loaders['train'].dataset.shuffle()

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer_rnn.py

@@ -0,0 +1,330 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import os.path as osp
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import wandb
+from rekognition_online_action_detection.evaluation import compute_result
+from tqdm import tqdm
+
+
+def log_results(phase, cfg, compute_result, gt_targets, pred_scores, losses, data_loader, epoch=None, best_mAP=None, best_epoch=None):
+    result = compute_result['perframe'](cfg, gt_targets, pred_scores)
+    per_class_AP = result['mean_AP']
+    
+    per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items())
+    
+    log = []
+    if epoch is not None:
+        log.append(f'\nEpoch {epoch:2}')
+    
+    log.append(f'{phase} det_loss: {losses["det"] / len(data_loader.dataset):.5f} det_mAP: {result["mean_AP"]:.5f} {per_class_AP_log}')
+    
+    if cfg.USE_WANDB:
+        wandb.log({
+            f"{phase.capitalize()} mAP": result['mean_AP'],
+            f"{phase.capitalize()} per_class_AP": per_class_AP,
+            f"{phase.capitalize()} Loss": losses['det'] / len(data_loader.dataset)
+        })
+    
+    if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+        for loss_type in ['verb', 'noun']:
+            if loss_type in losses:
+                log.append(f'{phase} {loss_type}_loss: {losses[loss_type] / len(data_loader.dataset):.5f}')
+    
+    return log, result['mean_AP']
+
+def do_perframe_det_train_rnn(cfg,
+                          data_loaders,
+                          model,
+                          criterion,
+                          optimizer,
+                          scheduler,
+                          device,
+                          checkpointer,
+                          logger):
+    # Setup model on multiple GPUs
+    if torch.cuda.device_count() > 1:
+        model = nn.DataParallel(model)
+
+    best_mAP, best_epoch = 0, 0
+    best_per_class_AP_log = None
+
+    for epoch in range(cfg.SOLVER.START_EPOCH, cfg.SOLVER.START_EPOCH + cfg.SOLVER.NUM_EPOCHS):
+        # Reset
+        det_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        pred_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES}
+        
+        det_pred_scores = []
+        det_gt_targets = []
+        aa_pred_scores = []
+        aa_gt_targets = []
+        
+        train_det_pred_scores = []
+        train_det_gt_targets = []
+        train_aa_pred_scores = []
+        train_aa_gt_targets = []
+
+        start = time.time()
+        
+        for phase in cfg.SOLVER.PHASES:
+            training = phase == 'train'
+            model.train(training)
+
+            with torch.set_grad_enabled(training):
+                pbar = tqdm(data_loaders[phase],
+                            desc='{}ing epoch {}'.format(phase.capitalize(), epoch))
+                
+                # __getitem in LSTR:
+                # return (fusion_visual_inputs, fusion_motion_inputs,
+                #       fusion_object_inputs, memory_key_padding_mask, target)
+                # __getitem__ in RNN:
+                # vid, rgb_input, flow_input, text_input, target, ant_target
+                for batch_idx, data in enumerate(pbar, start=1):
+                    batch_size = data[1].shape[0]                                        
+                    det_target = data[-2].to(device)
+                    aa_target = data[-1].to(device)
+                    
+                    loss_names = list(zip(*cfg.MODEL.CRITERIONS))[0]
+                    assert 'PRED_FUTURE' not in list(zip(*cfg.MODEL.CRITERIONS))[0]
+                    # if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]:
+                    #     det_score, feat_out, feat_ori = model(*[x.to(device) for x in data[:-1]])
+                    # else:
+                    #     det_score = model(*[x.to(device) for x in data[1:4]])
+                                        
+                    out_dict = model(*[x.to(device) for x in data[1:4]])
+                    # out_dict['logits']: [B, window_size, num_classes]
+                    # out_dict['anticipation_logits']: [B, window_size, ant_len, num_classes]
+                    
+                    # [workmem_len + ant_len, batch_size, self.d_model]: [50, 16, 1024]
+                    
+                    det_score = out_dict['logits']
+                    det_score = det_score
+                    if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0:
+                        aa_score = out_dict['anticipation_logits'] # [B, window_size, anticipation_length, num_class]
+                        aa_score = aa_score[:, -1, :, :].reshape(-1, cfg.DATA.NUM_CLASSES)
+                        
+                        aa_target = aa_target.reshape(-1, cfg.DATA.NUM_CLASSES)
+                        det_loss = criterion[loss_names[0]](aa_score, aa_target)
+                        anticipation_logits = out_dict['anticipation_logits']
+        
+                    else:
+                        det_score = det_score.reshape(-1, cfg.DATA.NUM_CLASSES)
+                        det_target = det_target.reshape(-1, cfg.DATA.NUM_CLASSES)
+                        det_loss = criterion[loss_names[0]](det_score, det_target)
+
+                        det_loss = criterion[loss_names[0]](det_score, det_target)
+                        
+                        
+                    # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                    #     aa_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                    #     # [16, 10, 3]
+                    #     aa_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                        
+                    # if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY:
+                    #     det_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                    #     det_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :]
+                    
+                    # det_score: [16, 50, 3], aa_score: [16, 10, 3]
+                    # det_score can be different when aa_score 
+                    
+                    det_losses[phase] += det_loss.item() * batch_size
+                    
+                    # Output log for current batch
+                    pbar.set_postfix({
+                        'lr': '{:.7f}'.format(scheduler.get_last_lr()[0]),
+                        'det_loss': '{:.5f}'.format(det_loss.item()),
+                    })
+
+                    if training:
+                        optimizer.zero_grad()
+                        loss = det_loss
+                        if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]:
+                            loss += pred_loss
+                        if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                            loss += noun_loss + verb_loss
+                        if loss.item() != 0:
+                            loss.backward()
+                            optimizer.step()
+                            scheduler.step()
+                        
+                        if cfg.USE_WANDB:
+                            wandb.log({"Learning Rate": scheduler.get_last_lr()[0]})
+                            wandb.log({"Train Loss": det_loss.item()})
+                        
+                        det_score = det_score.softmax(dim=1).cpu().tolist()
+                        det_target = det_target.cpu().tolist()
+                        train_det_pred_scores.extend(det_score)
+                        train_det_gt_targets.extend(det_target)
+                        # len(train_det_pred_scores) + 16
+                        
+                        if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0:
+                            aa_score = aa_score.softmax(dim=1).cpu().tolist()
+                            aa_target = aa_target.cpu().tolist()
+                            train_aa_pred_scores.extend(aa_score)
+                            train_aa_gt_targets.extend(aa_target)
+                            
+                    else:
+                        # Prepare for evaluation
+                        det_score = det_score.softmax(dim=1).cpu().tolist()
+                        det_target = det_target.cpu().tolist()
+                        det_pred_scores.extend(det_score)
+                        det_gt_targets.extend(det_target)
+                        
+                        if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0:
+                            aa_score = aa_score.softmax(dim=1).cpu().tolist()
+                            aa_target = aa_target.cpu().tolist()
+                            aa_pred_scores.extend(aa_score)
+                            aa_gt_targets.extend(aa_target)
+                        
+                        # print(f'len(det_score): {len(det_score)}, len(det_pred_scores): {len(det_pred_scores)}')  
+                        # print(f'len(aa_score): {len(aa_score)}, len(aa_pred_scores): {len(aa_pred_scores)}')
+            
+        end = time.time()
+
+        # After the phase loop, compute and log results
+        for phase in cfg.SOLVER.PHASES:
+            if phase == 'train':
+                gt_targets = train_det_gt_targets
+                pred_scores = train_det_pred_scores
+            else:  # 'test'
+                gt_targets = det_gt_targets
+                pred_scores = det_pred_scores
+            
+            result = compute_result['perframe'](cfg, gt_targets, pred_scores)
+            per_class_AP = result['per_class_AP']
+            
+            log = []
+            log.append(f'\nEpoch {epoch:2}')
+            log.append(f'{phase} det_loss: {det_losses[phase] / len(data_loaders[phase].dataset):.5f} det_mAP: {result["mean_AP"]:.5f}')
+            
+            per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items())
+            log.append(f'Per-class AP: {per_class_AP_log}')
+            
+            if cfg.USE_WANDB:
+                wandb.log({
+                    f"{phase.capitalize()} mAP": result['mean_AP'],
+                    f"{phase.capitalize()} per_class_AP": per_class_AP,
+                    f"{phase.capitalize()} Loss": det_losses[phase] / len(data_loaders[phase].dataset)
+                })
+            
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                log.append(f'{phase} verb_loss: {verb_losses[phase] / len(data_loaders[phase].dataset):.5f}')
+                log.append(f'{phase} noun_loss: {noun_losses[phase] / len(data_loaders[phase].dataset):.5f}')
+            
+            if phase == 'test':
+                if result['mean_AP'] > best_mAP:
+                    best_mAP = result['mean_AP']
+                    best_epoch = epoch
+                    best_per_class_AP_log = per_class_AP_log
+                    checkpointer.save_best(epoch, model, optimizer)
+                
+                log.append(f'\nbest epoch: {best_epoch:2} best det_mAP: {best_mAP:.5f} best per-class AP: {best_per_class_AP_log}\n')
+            
+            logger.info(' | '.join(log))
+
+        log.append(f'running time: {end - start:.2f} sec')
+        logger.info(' | '.join(log))
+
+        # Save checkpoint for model and optimizer
+        if cfg.SOLVER.SAVE_EVERY_EPOCHS > 0 and epoch % cfg.SOLVER.SAVE_EVERY_EPOCHS == 0:
+            checkpointer.save(epoch, model, optimizer)
+
+        # Shuffle dataset for next epoch
+        data_loaders['train'].dataset.shuffle()
+
+# @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()
+        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, text_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, text_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), text_input.cuda(), target.cuda(), ant_target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():
+                out_dict = model(rgb_input, flow_input, text_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, text_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/LSTR/src/rekognition_online_action_detection/engines/engines.py

@@ -0,0 +1,39 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from rekognition_online_action_detection.utils.registry import Registry
+
+TRAINERS = Registry()
+INFERENCES = Registry()
+
+def do_train(cfg,
+             data_loaders,
+             model,
+             criterion,
+             optimizer,
+             scheduler,
+             device,
+             checkpointer,
+             logger):
+    
+    return TRAINERS[cfg.MODEL.MODEL_NAME](
+        cfg,
+        data_loaders,
+        model,
+        criterion,
+        optimizer,
+        scheduler,
+        device,
+        checkpointer,
+        logger)
+
+
+def do_inference(cfg,
+                 model,
+                 device,
+                 logger):
+    return INFERENCES[cfg.MODEL.MODEL_NAME](
+        cfg,
+        model,
+        device,
+        logger)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_inference.py

@@ -0,0 +1,417 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os.path as osp
+from bisect import bisect_right
+import time
+import pickle as pkl
+
+import torch
+import numpy as np
+
+from rekognition_online_action_detection.evaluation import compute_result
+
+from ..base_inferences.perframe_det_batch_inference import do_perframe_det_batch_inference
+
+
+from ..engines import INFERENCES as registry
+
+
+@registry.register('LSTR')
+def do_lstr_batch_inference(cfg,
+                            model,
+                            device,
+                            logger):
+    if cfg.MODEL.LSTR.INFERENCE_MODE == 'stream':
+        do_lstr_stream_inference(cfg,
+                                 model,
+                                 device,
+                                 logger)
+    else:
+        do_perframe_det_batch_inference(cfg,
+                                        model,
+                                        device,
+                                        logger)
+
+
+def do_lstr_stream_inference(cfg, model, device, logger):
+    # Setup model to test mode
+    model.eval()
+
+    # Collect scores and targets
+    pred_scores = []
+    gt_targets = []
+
+    def to_device(x, dtype=np.float32):
+        return torch.as_tensor(x.astype(dtype)).unsqueeze(0).to(device)
+
+    long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH
+    long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE
+    long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES
+    work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH
+    work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE
+    work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES
+
+    # if len(cfg.DATA.TEST_SESSION_SET) != 1:
+    #     raise RuntimeError('Only support testing one video each time for stream inference, will fix later')
+
+    pred_scores_all = {}
+    gt_targets_all = {}
+
+    with torch.no_grad():
+        for session_idx, session in enumerate(cfg.DATA.TEST_SESSION_SET):
+            model.clear_cache()
+
+            if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                pred_scores = [[] for _ in range(cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)]
+            else:
+                pred_scores = []
+            gt_targets = []
+
+            visual_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.VISUAL_FEATURE, session + '.npy'), mmap_mode='r')
+            motion_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.MOTION_FEATURE, session + '.npy'), mmap_mode='r')
+            object_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.OBJECT_FEATURE, session + '.npy'), mmap_mode='r')
+            target = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.TARGET_PERFRAME, session + '.npy'))
+
+            start_time = time.time()
+
+            for work_start, work_end in zip(range(0, target.shape[0] + 1),
+                                            range(work_memory_length, target.shape[0] + 1)):
+                # Get target
+                # target = target[::work_memory_sample_rate]
+
+                # Get work memory
+                work_indices = np.arange(work_start, work_end).clip(0)
+                work_indices = work_indices[::work_memory_sample_rate]
+                work_visual_inputs = to_device(visual_inputs[work_indices])
+                work_motion_inputs = to_device(motion_inputs[work_indices])
+                work_object_inputs = to_device(object_inputs[work_indices])
+
+                # Get long memory
+                PRECISE = True
+                if PRECISE:
+                    long_end = work_start - long_memory_sample_rate
+                    if long_end < 0:
+                        long_indices = [0 for _ in range(long_memory_num_samples)]
+                        long_indices_set = [long_indices for _ in range(long_memory_sample_rate)]
+                        long_visual_inputs = to_device(visual_inputs[long_indices])
+                        long_motion_inputs = to_device(motion_inputs[long_indices])
+                        long_object_inputs = to_device(object_inputs[long_indices])
+                    else:
+                        long_indices = long_indices_set[long_end % long_memory_sample_rate][1:] + [long_end]
+                        long_indices_set[long_end % long_memory_sample_rate] = long_indices
+                        long_visual_inputs = to_device(visual_inputs[[long_end]])
+                        long_motion_inputs = to_device(motion_inputs[[long_end]])
+                        long_object_inputs = to_device(object_inputs[[long_end]])
+                else:
+                    long_end = work_start - 1
+                    if long_end == -1:
+                        long_indices = [0 for _ in range(long_memory_num_samples)]
+                        long_visual_inputs = to_device(visual_inputs[long_indices])
+                        long_motion_inputs = to_device(motion_inputs[long_indices])
+                        long_object_inputs = to_device(object_inputs[long_indices])
+                    elif long_end % long_memory_sample_rate == 0:
+                        long_indices = long_indices[1:] + [long_end]
+                        long_visual_inputs = to_device(visual_inputs[[long_end]])
+                        long_motion_inputs = to_device(motion_inputs[[long_end]])
+                        long_object_inputs = to_device(object_inputs[[long_end]])
+                    else:
+                        long_visual_inputs = None
+                        long_motion_inputs = None
+                        long_object_inputs = None
+
+                # Get memory key padding mask
+                memory_key_padding_mask = np.zeros(len(long_indices))
+                last_zero = bisect_right(long_indices, 0) - 1
+                if last_zero > 0:
+                    memory_key_padding_mask[:last_zero] = float('-inf')
+                memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32)).unsqueeze(0).to(device)
+
+                score = model.stream_inference(
+                    long_visual_inputs,
+                    long_motion_inputs,
+                    long_object_inputs,
+                    work_visual_inputs,
+                    work_motion_inputs,
+                    work_object_inputs,
+                    memory_key_padding_mask,
+                    cache_num=long_memory_sample_rate if PRECISE else 1,
+                    cache_id=long_end % long_memory_sample_rate if PRECISE else 0)[0]
+
+                if cfg.DATA.DATA_NAME.startswith('EK'):
+                    score = score.cpu().numpy()
+                else:
+                    score = score.softmax(dim=-1).cpu().numpy()
+
+                if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                    if work_start == 0:
+                        upsample_score = score[:-cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES].repeat(work_memory_sample_rate, axis=0)
+                        upsample_score = upsample_score[work_start:work_end-1]
+                        anticipate_score = score[-cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:]
+                        # coarse but sufficient
+                        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                            combined_score = np.concatenate((upsample_score, anticipate_score[:t_a + 1]),
+                                                            axis=0)
+                            pred_scores[t_a].extend(list(combined_score))
+                        gt_targets.extend(list(target[:work_end-1]))
+                    else:
+                        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                            pred_scores[t_a].append(list(score[t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES]))
+                        gt_targets.append(list(target[work_end - 1]))
+                else:
+                    if work_start == 0:
+                        gt_targets.extend(list(target[:work_end]))
+                        pred_scores.extend(list(score))
+                    else:
+                        gt_targets.append(list(target[work_end - 1]))
+                        pred_scores.append(list(score[-1]))
+
+            end_time = time.time()
+            logger.info('Running time: {:.3f} seconds'.format(end_time - start_time))
+
+            if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                    result = compute_result['perframe'](
+                        cfg,
+                        gt_targets,
+                        pred_scores[t_a][:-1 - t_a],
+                    )
+                    sec = (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE
+                    logger.info('mAP of video ({:.2f}s) {}: {:.5f}'.format(sec, session, result['mean_AP']))
+            else:
+                result = compute_result['perframe'](
+                    cfg,
+                    gt_targets,
+                    pred_scores,
+                )
+                logger.info('mAP of video {}: {:.5f}'.format(session, result['mean_AP']))
+
+
+            gt_targets_all[session] = np.array(gt_targets)
+            if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+                for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+                    pred_scores[t_a] = np.array(pred_scores[t_a][:- 1 - t_a])
+                pred_scores_all[session] = np.stack(pred_scores, axis=0).transpose(1, 2, 0)
+                pred_scores_all[session] = pred_scores_all[session]
+            else:
+                pred_scores_all[session] = np.array(pred_scores)
+
+
+    # pkl.dump({
+    #     'cfg': cfg,
+    #     'perframe_pred_scores': pred_scores_all,
+    #     'perframe_gt_targets': gt_targets_all,
+    # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.stream.pkl', 'wb'))
+
+    if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0:
+        maps_list = []
+        for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES):
+            result = compute_result['perframe'](
+                cfg,
+                np.concatenate(list(gt_targets_all.values()), axis=0),
+                np.concatenate(list(pred_scores_all.values()), axis=0)[:, :, t_a],
+            )
+            logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format(
+                (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+                cfg.DATA.METRICS, result['mean_AP']
+            ))
+            maps_list.append(result['mean_AP'])
+        logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format(
+            cfg.DATA.METRICS, np.mean(maps_list)
+        ))
+    else:
+        result = compute_result['perframe'](
+            cfg,
+            np.concatenate(list(gt_targets_all.values()), axis=0),
+            np.concatenate(list(pred_scores_all.values()), axis=0),
+        )
+        logger.info('Action detection perframe m{}: {:.5f}'.format(
+            cfg.DATA.METRICS, result['mean_AP']
+        ))
+
+    gt_targets = gt_targets_all
+    pred_scores = pred_scores_all
+    pred_scores_verb = {}
+    pred_scores_noun = {}
+
+    # segment-level evaluation on EK55/EK100
+    if cfg.DATA.DATA_NAME.startswith('EK'):
+        try:
+            import pandas as pd
+            from external.rulstm.RULSTM.utils import (get_marginal_indexes, marginalize, softmax,
+                                                      topk_accuracy_multiple_timesteps,
+                                                      topk_recall_multiple_timesteps,
+                                                      tta)
+        except:
+            raise ModuleNotFoundError
+        if cfg.DATA.DATA_NAME == 'EK55':
+            path_to_data = 'external/rulstm/RULSTM/data/ek55/'
+        elif cfg.DATA.DATA_NAME == 'EK100':
+            path_to_data = 'external/rulstm/RULSTM/data/ek100/'
+        segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'),
+                                   names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'],
+                                   skipinitialspace=True)
+        predictions, labels, ids = [], [], []
+        predictions_verb, predictions_noun = [], []
+        discarded_labels, discarded_ids = [], []
+        for segment in segment_list.iterrows():
+            start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS))
+            vid = segment[1]['video'].strip()
+            preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+            if cfg.MODEL.LSTR.V_N_CLASSIFIER:
+                raise NotImplementedError(
+                        'For TeSTra with decomposed classifier (V+N), please use batch mode, ie ``MODEL.LSTR.INFERENCE_MODE batch``'
+                )
+                preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+                preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1)
+            else:
+                preds_v = None
+                preds_n = None
+            if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH:
+                predictions.append(preds)
+                if preds_v is not None:
+                    predictions_verb.append(preds_v)
+                if preds_n is not None:
+                    predictions_noun.append(preds_n)
+                labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+                ids.append(segment[1]['id'])
+            else:
+                discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int))
+                discarded_ids.append(segment[1]['id'])
+        
+        actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id')
+
+        action_scores = np.concatenate(predictions)
+        if len(predictions_verb) > 0:
+            verb_scores = np.concatenate(predictions_verb)
+        else:
+            verb_scores = None
+        if len(predictions_noun) > 0:
+            noun_scores = np.concatenate(predictions_noun)
+        else:
+            noun_scores = None
+        labels = np.array(labels)
+        ids = np.array(ids)
+
+        vi = get_marginal_indexes(actions, 'verb')
+        ni = get_marginal_indexes(actions, 'noun')
+
+        action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1]))
+        if verb_scores is None:
+            verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0],
+                                                                action_scores.shape[1],
+                                                                -1)
+        else:
+            verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1]))
+            verb_scores = verb_scores.reshape(action_scores.shape[0],
+                                              action_scores.shape[1],
+                                              -1)
+        if noun_scores is None:
+            noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0],
+                                                                action_scores.shape[1],
+                                                                -1)
+        else:
+            noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1]))
+            noun_scores = noun_scores.reshape(action_scores.shape[0],
+                                              action_scores.shape[1],
+                                              -1)
+
+        include_discarded = (cfg.DATA.DATA_NAME == 'EK100')
+        if include_discarded:
+            dlab = np.array(discarded_labels)
+            dislab = np.array(discarded_ids)
+            ids = np.concatenate([ids, dislab])
+            num_disc = len(dlab)
+            labels = np.concatenate([labels, dlab])
+            verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:]))))
+            noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:]))))
+            action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:]))))
+
+        verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2]
+        verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels)
+        noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels)
+        action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels)
+
+        if cfg.DATA.DATA_NAME == 'EK55':
+            many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values
+            many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values
+            actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'))
+            a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values)
+                    for a in actions.iterrows()}
+            many_shot_actions = [a for a, (v, n) in a_to_vn.items()
+                                if v in many_shot_verbs or n in many_shot_nouns]
+            
+            verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels,
+                                                        k=5, classes=many_shot_verbs)
+            noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels,
+                                                        k=5, classes=many_shot_nouns)
+            action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels,
+                                                            k=5, classes=many_shot_actions)
+            
+            all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies,
+                                            verb_recalls, noun_recalls, action_recalls])
+            all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]]
+            indices = [
+                ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'),
+                ('Verb', 'Mean Top-5 Recall'),
+                ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'),
+                ('Noun', 'Mean Top-5 Recall'),
+                ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'),
+                ('Action', 'Mean Top-5 Recall'),
+            ]
+        else:
+            overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5)
+            overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5)
+            overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5)
+
+            unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True)
+            tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True)
+            tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True)
+            tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True)
+
+            unseen_bool_idx = pd.Series(ids).isin(unseen).values
+            tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values
+            tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values
+            tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values
+
+            tail_verb_recalls = topk_recall_multiple_timesteps(
+                verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5)
+            tail_noun_recalls = topk_recall_multiple_timesteps(
+                noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5)
+            tail_action_recalls = topk_recall_multiple_timesteps(
+                action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5)
+            
+            unseen_verb_recalls = topk_recall_multiple_timesteps(
+                verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5)
+            unseen_noun_recalls = topk_recall_multiple_timesteps(
+                noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5)
+            unseen_action_recalls = topk_recall_multiple_timesteps(
+                action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5)
+
+            all_accuracies = np.concatenate(
+                [overall_verb_recalls, overall_noun_recalls, overall_action_recalls,
+                 unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls,
+                 tail_verb_recalls, tail_noun_recalls, tail_action_recalls]
+            ) #9 x 8
+
+            indices = [
+                ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'),
+                ('Overall Mean Top-5 Recall', 'Action'),
+                ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'),
+                ('Unseen Mean Top-5 Recall', 'Action'),
+                ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'),
+                ('Tail Mean Top-5 Recall', 'Action'),
+            ]
+
+        cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH,
+                        1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE,
+                        cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str)
+        scores = pd.DataFrame(all_accuracies * 100, columns=cc,
+                            index=pd.MultiIndex.from_tuples(indices))
+        logger.info(scores)
+
+        tta_verb = tta(verb_scores, verb_labels)
+        tta_noun = tta(noun_scores, noun_labels)
+        tta_action = tta(action_scores, action_labels)
+        logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} '
+                    f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}')

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_trainer.py

@@ -0,0 +1,26 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from ..base_trainers.perframe_det_trainer import do_perframe_det_train
+from ..engines import TRAINERS as registry
+
+
+@registry.register('LSTR')
+def do_lstr_train(cfg,
+                  data_loaders,
+                  model,
+                  criterion,
+                  optimizer,
+                  scheduler,
+                  device,
+                  checkpointer,
+                  logger):
+    do_perframe_det_train(cfg,
+                          data_loaders,
+                          model,
+                          criterion,
+                          optimizer,
+                          scheduler,
+                          device,
+                          checkpointer,
+                          logger)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/__init__.py

@@ -0,0 +1,4 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .evalution import compute_result

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/evalution.py

@@ -0,0 +1,42 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from rekognition_online_action_detection.utils.registry import Registry
+
+from .metrics import perframe_average_precision, perstage_average_precision
+from .postprocessing import postprocessing as default_pp
+
+compute_result = Registry()
+
+
+@compute_result.register('perframe')
+def eval_perframe(cfg, ground_truth, prediction, **kwargs):
+    class_names = kwargs.get('class_names', cfg.DATA.CLASS_NAMES)
+    ignore_index = kwargs.get('ignore_index', cfg.DATA.IGNORE_INDEX)
+    metrics = kwargs.get('metrics', cfg.DATA.METRICS)
+    postprocessing = kwargs.get('postprocessing', default_pp(cfg.DATA.DATA_NAME))
+    
+    return perframe_average_precision(
+        ground_truth=ground_truth,
+        prediction=prediction,
+        class_names=class_names,
+        ignore_index=ignore_index,
+        metrics=metrics,
+        postprocessing=postprocessing,
+    )
+
+
+@compute_result.register('perstage')
+def eval_perstage(cfg, ground_truth, prediction, **kwargs):
+    class_names = kwargs.get('class_names', cfg.DATA.CLASS_NAMES)
+    ignore_index = kwargs.get('ignore_index', cfg.DATA.IGNORE_INDEX)
+    metrics = kwargs.get('metrics', cfg.DATA.METRICS)
+    postprocessing = kwargs.get('postprocessing', default_pp(cfg.DATA.DATA_NAME))
+    return perstage_average_precision(
+        ground_truth=ground_truth,
+        prediction=prediction,
+        class_names=class_names,
+        ignore_index=ignore_index,
+        metrics=metrics,
+        postprocessing=postprocessing,
+    )

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/metrics.py

@@ -0,0 +1,138 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+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(ground_truth,
+                               prediction,
+                               class_names,
+                               ignore_index,
+                               metrics,
+                               postprocessing):
+    """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':
+        compute_score = calibrated_average_precision_score
+    else:
+        raise RuntimeError('Unknown metrics: {}'.format(metrics))
+
+    # set ignore classes including background
+    ignore_set = set([ignore_index])
+    # enable background according to ignore_index
+    # ignore_set = set([ignore_index])
+
+    # Compute average precision
+    result['per_class_AP'] = OrderedDict()
+    
+    # when training, class_names becomes 2!
+    for idx, class_name in enumerate(class_names):
+        if idx not in ignore_set:
+            if np.any(ground_truth[:, idx]):
+                result['per_class_AP'][class_name] = compute_score(
+                    ground_truth[:, idx], prediction[:, idx])
+    result['mean_AP'] = np.mean(list(result['per_class_AP'].values()))
+    
+    return result
+
+
+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(ground_truth,
+                               prediction,
+                               class_names,
+                               ignore_index,
+                               metrics,
+                               postprocessing):
+    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))
+
+    # Ignore backgroud class
+    ignore_set = set([0, ignore_index])
+
+    # 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_set:
+                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/LSTR/src/rekognition_online_action_detection/evaluation/postprocessing.py

@@ -0,0 +1,36 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+
+def postprocessing(data_name):
+
+    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]
+
+    return {'THUMOS': thumos_postprocessing}.get(data_name, None)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/__init__.py

@@ -0,0 +1,5 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .lstr import *
+from .models import build_model

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/feature_head.py

@@ -0,0 +1,95 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['build_feature_head']
+
+import torch
+import torch.nn as nn
+from rekognition_online_action_detection.config.feature_sizes import \
+    FEATURE_SIZES
+from rekognition_online_action_detection.utils.registry import Registry
+
+FEATURE_HEADS = Registry()
+
+@FEATURE_HEADS.register('EasyCom')
+@FEATURE_HEADS.register('EGO4D')
+@FEATURE_HEADS.register('EGO4DVAL')
+@FEATURE_HEADS.register('THUMOS')
+@FEATURE_HEADS.register('TVSeries')
+@FEATURE_HEADS.register('EK55')
+@FEATURE_HEADS.register('EK100')
+@FEATURE_HEADS.register('YTConv')
+class BaseFeatureHead(nn.Module):
+    def __init__(self, cfg):
+        super(BaseFeatureHead, self).__init__()
+
+        if cfg.INPUT.MODALITY in ['visual', 'motion', 'object',
+                                  'visual+motion', 'motion+object', 'visual+object',
+                                  'visual+motion+object']:
+            self.with_visual = 'visual' in cfg.INPUT.MODALITY
+            self.with_motion = 'motion' in cfg.INPUT.MODALITY
+            self.with_object = 'object' in cfg.INPUT.MODALITY
+        else:
+            raise RuntimeError('Unknown modality of {}'.format(cfg.INPUT.MODALITY))
+        
+        INPUT_VISUAL_FEATURE = cfg.INPUT.VISUAL_FEATURE
+        INPUT_MOTION_FEATURE = cfg.INPUT.MOTION_FEATURE
+        INPUT_OBJECT_FEATURE = cfg.INPUT.OBJECT_FEATURE
+        
+        if self.with_visual and self.with_motion and self.with_object:
+            visual_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE]
+            motion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE]
+            object_size = FEATURE_SIZES[INPUT_OBJECT_FEATURE]
+            fusion_size = visual_size + motion_size + object_size
+        elif self.with_visual and self.with_motion:
+            visual_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE]
+            motion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE]
+            fusion_size = visual_size + motion_size
+        elif self.with_visual:
+            fusion_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE]
+        elif self.with_motion:
+            fusion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE]
+        elif self.with_object:
+            fusion_size = FEATURE_SIZES[INPUT_OBJECT_FEATURE]
+        
+        self.fusion_size = fusion_size
+        self.d_model = fusion_size
+        
+        if cfg.MODEL.FEATURE_HEAD.LINEAR_ENABLED:
+            if cfg.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES != -1:
+                self.d_model = cfg.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES
+            self.input_linear = nn.Sequential(
+                nn.Linear(fusion_size, self.d_model),
+                nn.ReLU(inplace=True),
+            )
+        else:
+            self.input_linear = nn.Identity()
+        if cfg.MODEL.FEATURE_HEAD.DROPOUT > 0:
+            self.dropout = nn.Dropout(cfg.MODEL.FEATURE_HEAD.DROPOUT)
+        else:
+            self.dropout = None
+
+    def forward(self, visual_input, motion_input, object_input):
+        if self.with_visual and self.with_motion and self.with_object:
+            fusion_input = torch.cat((visual_input, motion_input,
+                                      object_input), dim=-1)
+        elif self.with_visual and self.with_motion:
+            fusion_input = torch.cat((visual_input, motion_input), dim=-1)
+        elif self.with_visual:
+            fusion_input = visual_input
+        elif self.with_motion:
+            fusion_input = motion_input
+        elif self.with_object:
+            fusion_input = object_input
+        if self.dropout is not None:
+            fusion_input = self.dropout(fusion_input)
+        fusion_input = self.input_linear(fusion_input)
+        return fusion_input
+
+
+def build_feature_head(cfg):
+    data_names = ['YTConv', 'EGO4D']
+    data_name = next((name for name in data_names if name in cfg.DATA.DATA_NAME), cfg.DATA.DATA_NAME)
+    
+    feature_head = FEATURE_HEADS[data_name]
+    return feature_head(cfg)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/lstr.py

@@ -0,0 +1,424 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..utils.ek_utils import action_to_noun_map, action_to_verb_map
+from . import transformer as tr
+from .feature_head import build_feature_head
+from .models import META_ARCHITECTURES as registry
+from .normalized_linear import NormalizedLinear
+
+
+class LSTR(nn.Module):
+
+    def __init__(self, cfg):
+        super(LSTR, self).__init__()
+
+        # Build long feature heads
+        self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES
+        self.long_enabled = self.long_memory_num_samples > 0
+        if self.long_enabled:
+            self.feature_head_long = build_feature_head(cfg)
+
+        # Build work feature head
+        self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES
+        self.work_enabled = self.work_memory_num_samples > 0
+        if self.work_enabled:
+            self.feature_head_work = build_feature_head(cfg)
+
+        self.d_model = self.feature_head_work.d_model
+        self.num_heads = cfg.MODEL.LSTR.NUM_HEADS
+        self.dim_feedforward = cfg.MODEL.LSTR.DIM_FEEDFORWARD
+        self.dropout = cfg.MODEL.LSTR.DROPOUT
+        self.activation = cfg.MODEL.LSTR.ACTIVATION
+        self.num_classes = cfg.DATA.NUM_CLASSES
+
+        self.use_v_n = cfg.MODEL.LSTR.V_N_CLASSIFIER
+
+        self.encoder_attention_type = cfg.MODEL.LSTR.ENC_ATTENTION_TYPE
+        self.decay_alpha = cfg.MODEL.LSTR.ENC_ATTENTION_DECAY
+
+        self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES
+
+        # Build position encoding
+        self.pos_encoding = tr.PositionalEncoding(self.d_model, self.dropout)
+
+        self.long_memory_use_pe = cfg.MODEL.LSTR.LONG_MEMORY_USE_PE
+        self.work_memory_use_pe = cfg.MODEL.LSTR.WORK_MEMORY_USE_PE
+        self.include_work = cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK
+        self.include_work2 = cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2
+        if self.include_work and self.include_work2:
+            raise ValueError('cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK and'
+                             'cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2 should'
+                             'not be both True')
+
+        # Build LSTR encoder
+        if self.long_enabled:
+            self.enc_queries = nn.ModuleList()
+            self.enc_modules = nn.ModuleList()
+            for i, param in enumerate(cfg.MODEL.LSTR.ENC_MODULE):
+                if param[0] != -1:
+                    self.enc_queries.append(nn.Embedding(param[0], self.d_model))
+                    # in TeSTra, not like LSTR, gets atteyntion_type, decay_alpha as parameter
+                    # however, the parameter does not change anything from LSTR
+                    enc_layer = tr.TransformerDecoderLayer(
+                        self.d_model, self.num_heads, self.dim_feedforward,
+                        self.dropout, self.activation,
+                        attention_type=self.encoder_attention_type if i == 0 else 'dotproduct',
+                        decay_alpha=self.decay_alpha if i == 0 else 1.0)
+                    self.enc_modules.append(tr.TransformerDecoder(
+                        enc_layer, param[1], tr.layer_norm(self.d_model, param[2])))
+                else:
+                    self.enc_queries.append(None)
+                    enc_layer = tr.TransformerEncoderLayer(
+                        self.d_model, self.num_heads, self.dim_feedforward,
+                        self.dropout, self.activation,
+                        attention_type=self.encoder_attention_type)
+                    self.enc_modules.append(tr.TransformerEncoder(
+                        enc_layer, param[1], tr.layer_norm(self.d_model, param[2])))
+        else:
+            self.register_parameter('enc_queries', None)
+            self.register_parameter('enc_modules', None)
+
+        # Build LSTR decoder
+        if self.long_enabled:
+            param = cfg.MODEL.LSTR.DEC_MODULE
+            dec_layer = tr.TransformerDecoderLayer(
+                self.d_model, self.num_heads, self.dim_feedforward,
+                self.dropout, self.activation)
+            self.dec_modules = tr.TransformerDecoder(
+                dec_layer, param[1], tr.layer_norm(self.d_model, param[2]))
+        else:
+            param = cfg.MODEL.LSTR.DEC_MODULE
+            dec_layer = tr.TransformerEncoderLayer(
+                self.d_model, self.num_heads, self.dim_feedforward,
+                self.dropout, self.activation)
+            self.dec_modules = tr.TransformerEncoder(
+                dec_layer, param[1], tr.layer_norm(self.d_model, param[2]))
+
+        # Build decode token (for anticipation)
+        if self.anticipation_num_samples > 0:
+            self.dec_query = nn.Embedding(self.anticipation_num_samples, self.d_model)
+        else:
+            self.register_parameter('dec_query', None)
+
+        # Build classifier
+        if cfg.MODEL.LSTR.DROPOUT_CLS > 0:
+            self.dropout_cls = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS)
+            if self.use_v_n:
+                self.dropout_cls_v = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS)
+                self.dropout_cls_n = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS)
+        else:
+            self.dropout_cls = None
+        if cfg.MODEL.LSTR.FC_NORM:
+            self.classifier = NormalizedLinear(self.d_model, self.num_classes)
+        else:
+            self.classifier = nn.Linear(self.d_model, self.num_classes)
+
+        if self.use_v_n:
+            a_to_v = action_to_verb_map(cfg.DATA.EK_EXT_PATH,
+                                        action_offset=True,
+                                        verb_offset=True)
+            a_to_n = action_to_noun_map(cfg.DATA.EK_EXT_PATH,
+                                        action_offset=True,
+                                        noun_offset=True)
+            num_verbs = max(set(a_to_v.values())) + 1
+            num_nouns = max(set(a_to_n.values())) + 1
+            print('Number of verbs (inc. 0):', num_verbs)
+            print('Number of nouns (inc. 0):', num_nouns)
+            self.classifier_verb = nn.Linear(self.d_model, num_verbs)
+            self.classifier_noun = nn.Linear(self.d_model, num_nouns)
+
+        self.pred_future = 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]
+
+    def forward(self, visual_inputs, motion_inputs, object_inputs, memory_key_padding_mask=None):
+        if self.long_enabled:
+            # Compute long memories
+            if self.long_memory_use_pe:
+                long_memories = self.pos_encoding(self.feature_head_long(
+                    visual_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                    motion_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                    object_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                ).transpose(0, 1))
+                long_memories = long_memories.transpose(0, 1) ## transpose back
+            else:
+                long_memories = self.feature_head_long(
+                    visual_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                    motion_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                    object_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None],
+                )
+            batch_size = long_memories.shape[0]
+            # if long_memories.ndim == 5:
+            #     long_memories = self.pos_encoding_3d_spatial(long_memories)
+            long_memories = long_memories.view(batch_size, -1, self.d_model).transpose(0, 1)
+            # print(long_memories.shape)   # T(*H*W), B, C
+
+            if len(self.enc_modules) > 0:
+                enc_queries = [
+                    enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1)
+                    if enc_query is not None else None
+                    for enc_query in self.enc_queries
+                ]
+
+                if self.include_work2:
+                    memory_key_padding_mask = F.pad(memory_key_padding_mask,
+                                                    (0, self.work_memory_num_samples),
+                                                    'constant', 0)
+                # Encode long memories
+                if enc_queries[0] is not None:
+                    long_memories = self.enc_modules[0](enc_queries[0], long_memories,
+                                                        memory_key_padding_mask=memory_key_padding_mask)
+                else:
+                    long_memories = self.enc_modules[0](long_memories)
+                for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]):
+                    if enc_query is not None:
+                        long_memories = enc_module(enc_query, long_memories)
+                    else:
+                        long_memories = enc_module(long_memories)
+
+        # Concatenate memories
+        if self.long_enabled:
+            memory = long_memories
+
+        if self.work_enabled:
+            # Compute work memories
+            if visual_inputs.ndim == 5:
+                visual_inputs_avg = visual_inputs[:, self.long_memory_num_samples:].mean((2, 3))
+                motion_inputs_avg = motion_inputs[:, self.long_memory_num_samples:].mean((2, 3))
+                object_inputs_avg = object_inputs[:, self.long_memory_num_samples:].mean((2, 3))
+            else:
+                visual_inputs_avg = visual_inputs[:, self.long_memory_num_samples:]
+                motion_inputs_avg = motion_inputs[:, self.long_memory_num_samples:]
+                object_inputs_avg = object_inputs[:, self.long_memory_num_samples:]
+            work_memories_no_pe = self.feature_head_work(
+                visual_inputs_avg,
+                motion_inputs_avg,
+                object_inputs_avg
+                # object_inputs[:, self.long_memory_num_samples:],
+            ).transpose(0, 1)
+            if self.work_memory_use_pe:
+                work_memories = self.pos_encoding(work_memories_no_pe,
+                    padding=self.long_memory_num_samples if self.long_memory_use_pe else 0)
+            else:
+                work_memories = work_memories_no_pe
+            if self.dec_query is not None:
+                anticipate_memories = self.pos_encoding(
+                    self.dec_query.weight.unsqueeze(1).repeat(1, work_memories.shape[1], 1),
+                    padding=(self.long_memory_num_samples + self.work_memory_num_samples
+                             if self.long_memory_use_pe else self.work_memory_num_samples))
+                work_memories = torch.cat((work_memories, anticipate_memories), dim=0)
+            # work_memories = self.feature_head_work(
+            #     visual_inputs[:, self.long_memory_num_samples:],
+            #     motion_inputs[:, self.long_memory_num_samples:],
+            # ).transpose(0, 1)
+
+
+            # Build mask
+            mask = tr.generate_square_subsequent_mask(
+                work_memories.shape[0])
+            mask = mask.to(work_memories.device)
+            # Compute output
+            if self.long_enabled:
+                if self.include_work:
+                    extended_memory = torch.cat((memory, work_memories), dim=0)
+                    extended_mask = F.pad(mask,
+                                          (memory.shape[0], 0),
+                                          'constant', 0)
+                    extended_mask = extended_mask.to(extended_memory.device)
+                    output = self.dec_modules(
+                        work_memories,
+                        memory=extended_memory,
+                        tgt_mask=mask,
+                        memory_mask=extended_mask,
+                    )
+                else:
+                    output = self.dec_modules(
+                        work_memories,
+                        memory=memory,
+                        tgt_mask=mask,
+                    )
+            else:
+                output = self.dec_modules(
+                    work_memories,
+                    src_mask=mask,
+                )
+
+        full_orig_feats = self.feature_head_work(
+                visual_inputs_avg,
+                motion_inputs_avg,
+                object_inputs_avg
+                # object_inputs[:, self.long_memory_num_samples:]
+        ).transpose(0, 1)
+        all_outputs_decoded = output
+
+        # Compute classification score
+        if self.dropout_cls is not None:
+            output_a = self.dropout_cls(output)
+        else:
+            output_a = output
+        
+        # output_a
+        # [workmem_len + ant_len, batch_size, self.d_model]: [50, 16, 1024]
+        # score (before transpose) [len, batch_size, num_classes]
+        # self.classifier: nn.Linear(self.d_model, self.num_classes)
+        score = self.classifier(output_a)
+        score = score.transpose(0, 1)
+        
+        if self.use_v_n:
+            output_v = self.dropout_cls_v(output)
+            output_n = self.dropout_cls_n(output)
+            score_verb = self.classifier_verb(output_v).transpose(0, 1)
+            score_noun = self.classifier_noun(output_n).transpose(0, 1)
+            score = (score, score_verb, score_noun)
+        if self.pred_future:
+            return (score,
+                    all_outputs_decoded[1: self.work_memory_num_samples + 1, ...],
+                    full_orig_feats)
+        else:
+            return score
+
+
+@registry.register('LSTR')
+class LSTRStream(LSTR):
+
+    def __init__(self, cfg):
+        super(LSTRStream, self).__init__(cfg)
+
+        ############################
+        # Cache for stream inference
+        ############################
+        self.long_memories_cache = None
+        self.compressed_long_memories_cache = None
+
+    def stream_inference(self,
+                         long_visual_inputs,
+                         long_motion_inputs,
+                         long_object_inputs,
+                         work_visual_inputs,
+                         work_motion_inputs,
+                         work_object_inputs,
+                         memory_key_padding_mask=None,
+                         cache_num=1,
+                         cache_id=0):
+        assert self.long_enabled, 'Long-term memory cannot be empty for stream inference'
+        assert len(self.enc_modules) > 0, 'LSTR encoder cannot be disabled for stream inference'
+
+        if ((long_visual_inputs is not None) and
+            (long_motion_inputs is not None) and
+            (long_object_inputs is not None)):
+            # Compute long memories
+            long_memories = self.feature_head_long(
+                long_visual_inputs if not self.include_work2 else torch.cat((long_visual_inputs, work_visual_inputs), dim=1),
+                long_motion_inputs if not self.include_work2 else torch.cat((long_motion_inputs, work_motion_inputs), dim=1),
+                long_object_inputs if not self.include_work2 else torch.cat((long_object_inputs, work_object_inputs), dim=1),
+            ).transpose(0, 1)
+
+            if self.long_memories_cache is None:
+                self.long_memories_cache = [long_memories for _ in range(cache_num)]
+            else:
+                self.long_memories_cache[cache_id] = torch.cat((
+                    self.long_memories_cache[cache_id][1:], long_memories
+                ))
+
+            long_memories = self.long_memories_cache[cache_id]
+            pos = self.pos_encoding.pe[:self.long_memory_num_samples, :]
+
+            enc_queries = [
+                enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1)
+                if enc_query is not None else None
+                for enc_query in self.enc_queries
+            ]
+
+            if self.include_work2:
+                memory_key_padding_mask = F.pad(memory_key_padding_mask,
+                                                (0, self.work_memory_num_samples),
+                                                'constant', 0)
+            # Encode long memories
+            long_memories = self.enc_modules[0].stream_inference(enc_queries[0], long_memories, pos,
+                                                                 memory_key_padding_mask=memory_key_padding_mask,
+                                                                 cache_num=cache_num, cache_id=cache_id)
+            self.compressed_long_memories_cache  = long_memories
+            for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]):
+                if enc_query is not None:
+                    long_memories = enc_module(enc_query, long_memories)
+                else:
+                    long_memories = enc_module(long_memories)
+        else:
+            long_memories = self.compressed_long_memories_cache
+
+            enc_queries = [
+                enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1)
+                if enc_query is not None else None
+                for enc_query in self.enc_queries
+            ]
+
+            # Encode long memories
+            for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]):
+                if enc_query is not None:
+                    long_memories = enc_module(enc_query, long_memories)
+                else:
+                    long_memories = enc_module(long_memories)
+
+        # Concatenate memories
+        if self.long_enabled:
+            memory = long_memories
+
+        if self.work_enabled:
+            # Compute work memories
+            work_memories = self.pos_encoding(self.feature_head_work(
+                work_visual_inputs,
+                work_motion_inputs,
+                work_object_inputs,
+            ).transpose(0, 1),
+            padding=self.long_memory_num_samples if self.long_memory_use_pe else 0)
+            if self.dec_query is not None:
+                anticipate_memories = self.pos_encoding(
+                    self.dec_query.weight.unsqueeze(1).repeat(1, work_memories.shape[1], 1),
+                    padding=(self.long_memory_num_samples + self.work_memory_num_samples
+                             if self.long_memory_use_pe else  self.work_memory_num_samples))
+                work_memories = torch.cat((work_memories, anticipate_memories))
+
+            # Build mask
+            mask = tr.generate_square_subsequent_mask(
+                work_memories.shape[0])
+            mask = mask.to(work_memories.device)
+
+            # Compute output
+            if self.long_enabled:
+                if self.include_work:
+                    extended_memory = torch.cat((memory, work_memories), dim=0)
+                    extended_mask = F.pad(mask,
+                                          (memory.shape[0], 0),
+                                          'constant', 0)
+                    extended_mask = extended_mask.to(extended_memory.device)
+                    output = self.dec_modules(
+                        work_memories,
+                        memory=extended_memory,
+                        tgt_mask=mask,
+                        memory_mask=extended_mask,
+                    )
+                else:
+                    output = self.dec_modules(
+                        work_memories,
+                        memory=memory,
+                        tgt_mask=mask,
+                    )
+            else:
+                output = self.dec_modules(
+                    work_memories,
+                    src_mask=mask,
+                )
+
+        # Compute classification score
+        score = self.classifier(output)
+        import pdb; pdb.set_trace()
+
+        return score.transpose(0, 1)
+
+    def clear_cache(self):
+        self.long_memories_cache = None
+        self.compressed_long_memories_cache = None
+        self.enc_modules[0].clear_cache()

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/models.py

@@ -0,0 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['build_model']
+
+from rekognition_online_action_detection.utils.registry import Registry
+
+META_ARCHITECTURES = Registry()
+
+def build_model(cfg, device=None):
+    model = META_ARCHITECTURES[cfg.MODEL.MODEL_NAME](cfg)
+    from .weights_init import weights_init
+
+    model.apply(weights_init)
+    return model.to(device)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/normalized_linear.py

@@ -0,0 +1,57 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F
+
+
+class NormalizedLinear(nn.Module):
+    r"""Normalized linear layer which was proposed by Jiaqi Wang et. al. in 
+        their technical report for LVIS workshop at ECCV 2020.
+        https://arxiv.org/pdf/2008.10032.pdf
+    Args:
+        in_features: size of each input sample
+        out_features: size of each output sample
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+        tau: a temperature factor, set as 20 by default by following the paper.
+    """
+    __constants__ = ['in_features', 'out_features']
+    in_features: int
+    out_features: int
+    weight: Tensor
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True, tau: float = 20) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.Tensor(out_features, in_features))
+        if bias:
+            self.bias = Parameter(torch.Tensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+        self.tau = tau
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: Tensor) -> Tensor:
+
+        return F.linear(
+            input / torch.norm(input, dim=-1, keepdim=True),
+            self.tau * self.weight /
+            torch.norm(self.weight, dim=-1, keepdim=True),
+            self.bias
+        )
+
+    def extra_repr(self) -> str:
+        return 'in_features={}, out_features={}, bias={}'.format(
+            self.in_features, self.out_features, self.bias is not None
+        )

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/__init__.py

@@ -0,0 +1,10 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .position_encoding import PositionalEncoding
+from .position_encoding import PositionalEncoding2D
+from .transformer import MultiheadAttention
+from .transformer import Transformer
+from .transformer import TransformerEncoder, TransformerEncoderLayer
+from .transformer import TransformerDecoder, TransformerDecoderLayer
+from .utils import layer_norm, generate_square_subsequent_mask

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/multihead_attention.py

@@ -0,0 +1,456 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch._C import device
+import torch.nn as nn
+import torch.nn.functional as F
+
+import time
+
+class DotProductAttention(nn.Module):
+
+    def __init__(self, dropout=0.0, alpha=1.0, num_heads=8):
+        super(DotProductAttention, self).__init__()
+
+        self.dropout = dropout
+        self.alpha = alpha
+        self.num_heads = num_heads
+        self.decay_mat = None
+
+    def forward(self, q, k, v, attn_mask=None):
+        attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+
+        if attn_mask is not None:
+            attn_output_weights += attn_mask
+
+        # attn_output_weights = F.softmax(attn_output_weights, dim=-1)
+        # attn_output_weights = F.dropout(attn_output_weights,
+        #                                 p=self.dropout,
+        #                                 training=self.training)
+        if self.alpha != 1.0 and self.decay_mat is None:
+            self.decay_mat = torch.vander(torch.tensor([self.alpha] * self.num_heads), N=v.shape[1])
+            self.decay_mat = self.decay_mat.to(v.device)
+        if self.alpha != 1.0:
+            bs = attn_output_weights.shape[0] // self.num_heads
+            decay_mat = self.decay_mat.unsqueeze(0).repeat(bs, 1, 1)
+            decay_mat = decay_mat.view(-1, *decay_mat.shape[2:])
+            attn_output_weights = torch.log(decay_mat[:, None, :] + 1e-10) + attn_output_weights
+        attn_output_weights = F.softmax(attn_output_weights, dim=-1)
+        attn_output_weights = F.dropout(attn_output_weights,
+                                        p=self.dropout,
+                                        training=self.training)
+        attn_output = torch.bmm(attn_output_weights, v)
+
+        return attn_output
+
+
+class DotProductAttentionStream(DotProductAttention):
+
+    def __init__(self, dropout=0.0, alpha=1.0):
+        super(DotProductAttentionStream, self).__init__(dropout, alpha)
+
+        ############################
+        # Cache for stream inference
+        ############################
+        self.k_weights_cache = None
+        self.k_pos_weights_cache = None
+        self.QK_exp_cache = None
+        self.recache_cnt = 0
+        self.recache_freq = 99999
+        self.alpha = alpha
+
+    def clear_cache(self):
+        self.k_weights_cache = None
+        self.k_pos_weights_cache = None
+        self.QK_exp_cache = None
+        self.recache_cnt = 0
+
+
+    def stream_inference(self, q, k, v, k_pos, v_pos, attn_mask=None, cache_num=1, cache_id=0):
+        '''
+        if self.k_weights_cache is not None:
+            k_weights_new = torch.bmm(q, k[:, [-1]].transpose(1, 2))
+            k_weights = torch.cat((self.k_weights_cache[:, :, 1:], k_weights_new), dim=-1)
+            self.k_weights_cache = k_weights
+            # k_pos_weights = self.k_pos_weights_cache
+        else:
+            k_weights = torch.bmm(q, k.transpose(1, 2))
+            self.k_weights_cache = k_weights
+            # k_pos_weights = torch.bmm(q, k_pos.transpose(1, 2))
+            # self.k_pos_weights_cache = k_pos_weights
+        # attn_output_weights = k_weights + k_pos_weights
+        attn_output_weights = k_weights
+
+        if attn_mask is not None:
+            attn_output_weights += attn_mask
+
+        attn_output_weights = F.softmax(attn_output_weights, dim=-1)
+        attn_output_weights = F.dropout(attn_output_weights,
+                                        p=self.dropout,
+                                        training=self.training)
+        # attn_output = torch.bmm(attn_output_weights, (v + v_pos))
+        attn_output = torch.bmm(attn_output_weights, v)
+        return attn_output
+        '''
+        B = 1
+        num = 1
+        q = q.repeat(B, 1, 1)
+        k = k.repeat(B, num, 1)
+        v = v.repeat(B, num, 1)
+        attn_mask = attn_mask.repeat(B, 1, num)
+        if self.alpha <= 1.0:
+            if self.QK_exp_cache is None:
+                QK = torch.bmm(q, k.transpose(1, 2))  # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C)
+                self.QK_max = QK.max(dim=-1, keepdim=True).values
+                QK_exp = torch.exp(QK - self.QK_max + attn_mask)  # => O(H*Tq*Tk)
+                self.QK_exp_cache = [QK_exp for _ in range(cache_num)]
+                self.QKV_exp_cache = [torch.bmm(QK_exp, v) for _ in range(cache_num)]       # => O(H*Tq*Tk*C)
+                self.Z = [QK_exp.sum(dim=-1, keepdim=True) for _ in range(cache_num)]       # => O(H*Tq*Tk)
+                output = self.QKV_exp_cache[cache_id] / self.Z[cache_id]
+                self.recache_cnt += 1
+                return output
+            else:
+                QK_push = torch.bmm(q, k[:, -num:, :].transpose(1, 2)) # => O(H*Tq*C)
+                QK_exp_push = torch.exp(QK_push - self.QK_max)        # => O(H*Tq*Tk)
+                sum_push = QK_exp_push.sum(-1, keepdim=True)
+                QKV_exp_push = torch.bmm(QK_exp_push, v[:, -num:, :])  # => O(H*Tq*Tk)
+                self.QKV_exp_cache[cache_id] = self.alpha * self.QKV_exp_cache[cache_id] + QKV_exp_push # => O(H*Tq*Tk)
+                self.Z[cache_id] = self.alpha * self.Z[cache_id] + sum_push
+
+                output = self.QKV_exp_cache[cache_id] / self.Z[cache_id]
+                self.recache_cnt += 1
+                return output
+
+        if self.QK_exp_cache is None:
+            QK = torch.bmm(q, k.transpose(1, 2))  # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C)
+            self.QK_max = QK.max(dim=-1, keepdim=True).values
+            QK_exp = torch.exp(QK - self.QK_max + attn_mask)  # => O(H*Tq*Tk)
+            self.QK_exp_cache = [QK_exp for _ in range(cache_num)]
+            self.QKV_exp_cache = [torch.bmm(QK_exp, v) for _ in range(cache_num)]       # => O(H*Tq*Tk*C)
+
+            self.Z = [QK_exp.sum(dim=-1, keepdim=True) for _ in range(cache_num)]       # => O(H*Tq*Tk)
+            self.QKV_exp_pop = [torch.bmm(self.QK_exp_cache[cache_i][:, :, :num], v[:, :num, :]) for cache_i in range(cache_num)]
+            self.sum_pop = [self.QK_exp_cache[cache_i][:, :, :num].sum(-1, keepdim=True) for cache_i in range(cache_num)]
+
+            output = self.QKV_exp_cache[cache_id] / self.Z[cache_id]
+            self.recache_cnt += 1
+            return output
+        elif (self.recache_cnt // cache_num) % self.recache_freq == 0:
+            QK = torch.bmm(q, k.transpose(1, 2))  # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C)
+            self.QK_max = QK.max(dim=-1, keepdim=True).values
+            QK_exp = torch.exp(QK - self.QK_max + attn_mask)  # => O(H*Tq*Tk)
+            self.QK_exp_cache[cache_id] = QK_exp
+            self.QKV_exp_cache[cache_id] = torch.bmm(QK_exp, v)         # => O(H*Tq*Tk*C)
+
+            self.Z[cache_id] = QK_exp.sum(dim=-1, keepdim=True)         # => O(H*Tq*Tk)
+
+            self.QKV_exp_pop[cache_id] = torch.bmm(self.QK_exp_cache[cache_id][:, :, :num], v[:, :num, :])
+            self.sum_pop[cache_id] = self.QK_exp_cache[cache_id][:, :, :num].sum(-1, keepdim=True)
+            output = self.QKV_exp_cache[cache_id] / self.Z[cache_id]
+            self.recache_cnt += 1
+            return output
+        else:
+            QK_push = torch.bmm(q, k[:, -num:, :].transpose(1, 2)) # => O(H*Tq*C)
+            QK_exp_push = torch.exp(QK_push - self.QK_max)        # => O(H*Tq*Tk)
+            sum_push = QK_exp_push.sum(-1, keepdim=True)
+            QKV_exp_push = torch.bmm(QK_exp_push, v[:, -num:, :])  # => O(H*Tq*Tk)
+
+            # add with no decay
+            self.QKV_exp_cache[cache_id] = self.QKV_exp_cache[cache_id] + QKV_exp_push # => O(H*Tq*Tk)
+
+            # add with no decay
+            self.Z[cache_id] = self.Z[cache_id] + sum_push  # => O(H*Tq*Tk)
+
+            self.QK_exp_cache[cache_id] = torch.cat((self.QK_exp_cache[cache_id][:, :, num:], QK_exp_push), dim=-1)
+            self.QKV_exp_pop[cache_id] = torch.bmm(self.QK_exp_cache[cache_id][:, :, :num], v[:, :num, :])
+            self.sum_pop[cache_id] = self.QK_exp_cache[cache_id][:, :, :num].sum(-1, keepdim=True)
+            output = self.QKV_exp_cache[cache_id] / self.Z[cache_id]
+            self.recache_cnt += 1
+            return output
+
+
+class MultiheadAttention(nn.Module):
+
+    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, kdim=None, vdim=None,
+                 attention_type='dotproduct', decay_alpha=1.0):
+        super(MultiheadAttention, self).__init__()
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        if self._qkv_same_embed_dim:
+            self.in_proj_weight = nn.Parameter(torch.empty(3 * embed_dim, embed_dim))
+        else:
+            raise RuntimeError('Do not support q, k, v have different dimensions')
+
+        if bias:
+            self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+        if self._qkv_same_embed_dim:
+            nn.init.xavier_uniform_(self.in_proj_weight)
+
+        if self.in_proj_bias is not None:
+            nn.init.constant_(self.in_proj_bias, 0.)
+            nn.init.constant_(self.out_proj.bias, 0.)
+
+        if attention_type == 'dotproduct':
+            self.attention = DotProductAttention(dropout, decay_alpha, num_heads)
+        elif attention_type == 'linear':
+            self.attention = LinearAttention(dropout)
+        else:
+            raise RuntimeError('attention_type should be [dotproduct | linear]')
+
+    def forward(self, q, k, v, attn_mask=None, key_padding_mask=None):
+        tsz, bsz, embed_dim = q.shape[0], q.shape[1], q.shape[2]
+
+        head_dim = embed_dim // self.num_heads
+        assert head_dim * self.num_heads == embed_dim, \
+            'embed_dim must be divisible by num_heads'
+        scaling = float(head_dim) ** -0.5
+
+        _b = self.in_proj_bias
+        _start = None
+        _end = embed_dim
+        _w = self.in_proj_weight[:_end, :]
+        if _b is not None:
+            _b = _b[:_end]
+        q = F.linear(q, _w, _b)
+
+        _b = self.in_proj_bias
+        _start = embed_dim
+        _end = embed_dim * 2
+        _w = self.in_proj_weight[_start:_end, :]
+        if _b is not None:
+            _b = _b[_start:_end]
+        k = F.linear(k, _w, _b)
+
+        _b = self.in_proj_bias
+        _start = embed_dim * 2
+        _end = None
+        _w = self.in_proj_weight[_start:, :]
+        if _b is not None:
+            _b = _b[_start:]
+        v = F.linear(v, _w, _b)
+
+        q = q * scaling
+
+        q = q.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0).repeat(bsz, 1, 1)
+            attn_mask = attn_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)
+            attn_mask = attn_mask.reshape(-1, *attn_mask.shape[2:])
+
+        if key_padding_mask is not None:
+            if key_padding_mask.ndim == 4:
+                key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1, 1, 1)
+                key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1, 1, 1)
+                key_padding_mask = key_padding_mask.reshape(*key_padding_mask.shape[:3], -1)
+            else:
+                key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1)
+                key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)
+            key_padding_mask = key_padding_mask.reshape(-1, *key_padding_mask.shape[2:])
+
+        if attn_mask is not None and key_padding_mask is not None:
+            mask = attn_mask + key_padding_mask
+        elif attn_mask is not None:
+            mask = attn_mask
+        elif key_padding_mask is not None:
+            mask = key_padding_mask
+        else:
+            mask = None
+
+        attn_output = self.attention(q, k, v, mask)
+        attn_output = attn_output.transpose(0, 1).contiguous().view(tsz, bsz,
+                                                                    self.embed_dim)
+        return self.out_proj(attn_output), None
+
+
+class MultiheadAttentionStream(MultiheadAttention):
+
+    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, kdim=None, vdim=None,
+                 attention_type='dotproduct', decay_alpha=1.0):
+        super(MultiheadAttentionStream, self).__init__(embed_dim, num_heads,
+                                                       dropout, bias, kdim,
+                                                       vdim, attention_type,
+                                                       decay_alpha)
+
+        if attention_type == 'dotproduct':
+            self.attention = DotProductAttentionStream(dropout, decay_alpha)
+        elif attention_type == 'linear':
+            self.attention = LinearAttentionStream(dropout)
+        else:
+            raise RuntimeError('attention_type should be [dotproduct | linear]')
+
+        ############################
+        # Cache for stream inference
+        ############################
+        self.q_cache = None
+        self.k_cache = None
+        self.v_cache = None
+        self.k_pos_cache = None
+        self.v_pos_cache = None
+
+    def clear_cache(self):
+        self.q_cache = None
+        self.k_cache = None
+        self.v_cache = None
+        self.k_pos_cache = None
+        self.v_pos_cache = None
+        self.attention.clear_cache()
+
+    def stream_inference(self, q, k, v, pos, attn_mask=None, key_padding_mask=None,
+                         cache_num=1, cache_id=0):
+        tsz, bsz, embed_dim = q.shape[0], q.shape[1], q.shape[2]
+
+        head_dim = embed_dim // self.num_heads
+        assert head_dim * self.num_heads == embed_dim, \
+            'embed_dim must be divisible by num_heads'
+        scaling = float(head_dim) ** -0.5
+
+        if self.q_cache is not None:
+            q = self.q_cache
+        else:
+            _b = self.in_proj_bias
+            _start = None
+            _end = embed_dim
+            _w = self.in_proj_weight[:_end, :]
+            if _b is not None:
+                _b = _b[:_end]
+            q = F.linear(q, _w, _b)
+            self.q_cache = q
+
+        assert (self.k_cache is None) == (self.k_pos_cache is None)
+        if self.k_cache is not None:
+            _b = self.in_proj_bias
+            _start = embed_dim
+            _end = embed_dim * 2
+            _w = self.in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            k_new = F.linear(k[[-1]], _w, None)
+            k = torch.cat((self.k_cache[1:], k_new))
+            self.k_cache = k
+            k_pos = self.k_pos_cache
+        else:
+            _b = self.in_proj_bias
+            _start = embed_dim
+            _end = embed_dim * 2
+            _w = self.in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            k = F.linear(k, _w, None)
+            self.k_cache = k
+            k_pos = F.linear(pos, _w, _b)
+            self.k_pos_cache = k_pos
+
+        assert (self.v_cache is None) == (self.v_pos_cache is None)
+        if self.v_cache is not None:
+            _b = self.in_proj_bias
+            _start = embed_dim * 2
+            _end = None
+            _w = self.in_proj_weight[_start:, :]
+            if _b is not None:
+                _b = _b[_start:]
+            v_new = F.linear(v[[-1]], _w, None)
+            v = torch.cat((self.v_cache[1:], v_new))
+            self.v_cache = v
+            v_pos = self.v_pos_cache
+        else:
+            _b = self.in_proj_bias
+            _start = embed_dim * 2
+            _end = None
+            _w = self.in_proj_weight[_start:, :]
+            if _b is not None:
+                _b = _b[_start:]
+            v = F.linear(v, _w, None)
+            self.v_cache = v
+            v_pos = F.linear(pos, _w, _b)
+            self.v_pos_cache = v_pos
+
+        q = q * scaling
+
+        q = q.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        k_pos = k_pos.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+        v_pos = v_pos.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0).repeat(bsz, 1, 1)
+            attn_mask = attn_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)
+            attn_mask = attn_mask.reshape(-1, *attn_mask.shape[2:])
+
+        if key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1)
+            key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)
+            key_padding_mask = key_padding_mask.reshape(-1, *key_padding_mask.shape[2:])
+
+        if attn_mask is not None and key_padding_mask is not None:
+            mask = attn_mask + key_padding_mask
+        elif attn_mask is not None:
+            mask = attn_mask
+        elif key_padding_mask is not None:
+            mask = key_padding_mask
+        else:
+            mask = None
+
+        attn_output = self.attention.stream_inference(q, k, v, k_pos, v_pos, mask,
+                                                      cache_num=cache_num, cache_id=cache_id)
+        attn_output = attn_output.transpose(0, 1).contiguous().view(tsz, bsz,
+                                                                    self.embed_dim)
+        return self.out_proj(attn_output), None
+
+
+class LinearAttention(nn.Module):
+
+    def __init__(self, dropout=0.0, nonlinear=F.elu,
+                 eps=1e-6):
+        super(LinearAttention, self).__init__()
+
+        self.dropout = dropout
+        self.nonlinear = nonlinear
+        self.eps = eps
+
+    def forward(self, q, k, v, attn_mask=None):
+        kk = self.nonlinear(k) + 1
+        qq = self.nonlinear(q) + 1
+        vv = self.nonlinear(v) + 1
+
+        kv = torch.bmm(kk.transpose(1, 2), vv)
+        
+        ks = torch.sum(kk, dim=1, keepdim=True)  # (B 1 C)
+        Z = 1 / (torch.bmm(qq, ks.transpose(1, 2)) + self.eps)
+        attn_output = torch.bmm(qq, kv.transpose(1, 2)) * Z
+
+        # if attn_mask is not None:
+        #     print(attn_mask.shape)
+        #     raise RuntimeError(("LinearAttention does not support arbitrary "
+        #                         "attention masks"))
+
+        return attn_output
+
+class LinearAttentionStream(LinearAttention):
+
+    def __init__(self, dropout=0.0):
+        super(LinearAttentionStream, self).__init__(dropout)
+
+        ############################
+        # Cache for stream inference
+        ############################
+        self.k_weights_cache = None
+        self.k_pos_weights_cache = None
+
+    def stream_inference(self, q, k, v, k_pos, v_pos, attn_mask=None):
+        raise NotImplementedError
+

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/position_encoding.py

@@ -0,0 +1,53 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        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, padding=0):
+        x = x + self.pe[padding: padding + x.shape[0], :]
+        return self.dropout(x)
+
+
+class PositionalEncoding2D(nn.Module):
+
+    def __init__(self, d_model, dropout=0.1, height=64, width=64):
+        super(PositionalEncoding2D, self).__init__()
+
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(height, width, d_model)
+        position_x = torch.arange(0, width, dtype=torch.float).unsqueeze(1)
+        position_y = torch.arange(0, height, dtype=torch.float).unsqueeze(1)
+        d_model = int(d_model / 2)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, :, 0:d_model:2] = torch.sin(position_x * div_term).unsqueeze(0).repeat(height, 1, 1)
+        pe[:, :, 1:d_model:2] = torch.cos(position_x * div_term).unsqueeze(0).repeat(height, 1, 1)
+        pe[:, :, d_model::2] = torch.sin(position_y * div_term).unsqueeze(1).repeat(1, width, 1)
+        pe[:, :, d_model + 1::2] = torch.cos(position_y * div_term).unsqueeze(1).repeat(1, width, 1)
+        pe = pe.unsqueeze(0).unsqueeze(0)    #.transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x, padding=(0,0)):
+        x = x + self.pe[:, :,
+                        padding[0]: padding[0] + x.shape[2],
+                        padding[1]: padding[1] + x.shape[3], :]
+        return self.dropout(x)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/transformer.py

@@ -0,0 +1,250 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .multihead_attention import MultiheadAttentionStream as MultiheadAttention
+
+
+class Transformer(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation='relu', custom_encoder=None, custom_decoder=None):
+        super(Transformer, self).__init__()
+
+        if custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
+            encoder_norm = nn.LayerNorm(d_model)
+            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
+            decoder_norm = nn.LayerNorm(d_model)
+            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
+                memory_mask=None, src_key_padding_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        if src.size(1) != tgt.size(1):
+            raise RuntimeError('the batch number of src and tgt must be equal')
+
+        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
+            raise RuntimeError('the feature number of src and tgt must be equal to d_model')
+
+        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
+        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
+                              tgt_key_padding_mask=tgt_key_padding_mask,
+                              memory_key_padding_mask=memory_key_padding_mask)
+        return output
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=src_mask,
+                         src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None):
+        super(TransformerDecoder, self).__init__()
+
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def clear_cache(self):
+        if len(self.layers) != 1:
+            raise RuntimeError('Number of layers cannot larger than 1 for stream inference')
+
+        self.layers[0].clear_cache()
+
+    def stream_inference(self, tgt, memory, pos, tgt_mask=None,
+                         memory_mask=None, tgt_key_padding_mask=None,
+                         memory_key_padding_mask=None, cache_num=1, cache_id=0):
+        output = tgt
+
+        if len(self.layers) != 1:
+            raise RuntimeError('Number of layers cannot larger than 1 for stream inference')
+
+        output = self.layers[0].stream_inference(output, memory, pos,
+                                                 tgt_mask=tgt_mask,
+                                                 memory_mask=memory_mask,
+                                                 tgt_key_padding_mask=tgt_key_padding_mask,
+                                                 memory_key_padding_mask=memory_key_padding_mask,
+                                                 cache_num=cache_num, cache_id=cache_id)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+    def forward(self, tgt, memory, tgt_mask=None,
+                memory_mask=None, tgt_key_padding_mask=None,
+                memory_key_padding_mask=None):
+        output = tgt
+
+        for mod in self.layers:
+            output = mod(output, memory, tgt_mask=tgt_mask,
+                         memory_mask=memory_mask,
+                         tgt_key_padding_mask=tgt_key_padding_mask,
+                         memory_key_padding_mask=memory_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerEncoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu'):
+        super(TransformerEncoderLayer, self).__init__()
+
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerEncoderLayer, self).__setstate__(state)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation='relu', attention_type='dotproduct',
+                 decay_alpha=1.0):
+        super(TransformerDecoderLayer, self).__init__()
+
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout,
+                                            attention_type=attention_type)
+        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout,
+                                                 attention_type=attention_type,
+                                                 decay_alpha=decay_alpha)
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+        ############################
+        # Cache for stream inference
+        ############################
+        self.tgt_cache = None
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerDecoderLayer, self).__setstate__(state)
+
+    def clear_cache(self):
+        self.tgt_cache = None
+        self.multihead_attn.clear_cache()
+
+    def stream_inference(self, tgt, memory, pos, tgt_mask=None, memory_mask=None,
+                         tgt_key_padding_mask=None, memory_key_padding_mask=None,
+                         cache_num=1, cache_id=0):
+        if self.tgt_cache is None:
+            tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
+                                  key_padding_mask=tgt_key_padding_mask)[0]
+            tgt = tgt + self.dropout1(tgt2)
+            tgt = self.norm1(tgt)
+            self.tgt_cache = tgt
+        else:
+            tgt = self.tgt_cache
+        tgt2 = self.multihead_attn.stream_inference(tgt, memory, memory, pos, attn_mask=memory_mask,
+                                                    key_padding_mask=memory_key_padding_mask,
+                                                    cache_num=cache_num, cache_id=cache_id)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    if activation == 'relu':
+        return F.relu
+    elif activation == 'gelu':
+        return F.gelu
+
+    raise RuntimeError('activation should be relu/gelu, not {}'.format(activation))

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/utils.py

@@ -0,0 +1,16 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+assert torch.__version__ >= '1.6.0'
+import torch.nn as nn
+
+
+def layer_norm(d_model, condition=True):
+    return nn.LayerNorm(d_model) if condition else None
+
+
+def generate_square_subsequent_mask(sz):
+    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+    return mask

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/weights_init.py

@@ -0,0 +1,33 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+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)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/__init__.py

@@ -0,0 +1,5 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .optimizers import build_optimizer
+from .lr_scheduler import build_scheduler

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/lr_scheduler.py

@@ -0,0 +1,224 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from bisect import bisect_left
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+def _get_warmup_factor_at_iter(warmup_method,
+                               this_iter,
+                               warmup_iters,
+                               warmup_factor):
+    if this_iter >= warmup_iters:
+        return 1.0
+    
+    if warmup_method == 'constant':
+        return warmup_factor
+    elif warmup_method == 'linear':
+        alpha = this_iter / warmup_iters
+        return warmup_factor * (1 - alpha) + alpha
+    else:
+        raise ValueError('Unknown warmup method: {}'.format(warmup_method))
+
+
+class MultiStepLR(_LRScheduler):
+
+    def __init__(self,
+                 optimizer,
+                 milestones,
+                 gamma=0.1,
+                 last_epoch=-1):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                'Milestones should be a list of increasing integers. Got {}'.format(milestones)
+            )
+        self.milestones = milestones
+        self.gamma = gamma
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        return [
+            base_lr
+            * self.gamma ** bisect_left(self.milestones, self.last_epoch)
+            for base_lr in self.base_lrs
+        ]
+
+    def _compute_values(self):
+        # The new interface
+        return self.get_lr()
+
+
+class WarmupMultiStepLR(_LRScheduler):
+
+    def __init__(self,
+                 optimizer,
+                 milestones,
+                 gamma=0.1,
+                 warmup_factor=0.3,
+                 warmup_iters=500,
+                 warmup_method='linear',
+                 last_epoch=-1):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                'Milestones should be a list of increasing integers. Got {}'.format(milestones)
+            )
+        self.milestones = milestones
+        self.gamma = gamma
+        self.warmup_factor = warmup_factor
+        self.warmup_iters = warmup_iters
+        self.warmup_method = warmup_method
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        warmup_factor = _get_warmup_factor_at_iter(
+            self.warmup_method,
+            self.last_epoch,
+            self.warmup_iters,
+            self.warmup_factor,
+        )
+        import pdb; pdb.set_trace()
+        return [
+            base_lr
+            * warmup_factor
+            * self.gamma ** bisect_left(self.milestones, self.last_epoch)
+            for base_lr in self.base_lrs
+        ]
+
+    def _compute_values(self):
+        # The new interface
+        return self.get_lr()
+
+
+class CosineLR(_LRScheduler):
+
+    def __init__(self,
+                 optimizer,
+                 max_iters,
+                 last_epoch=-1):
+        self.max_iters = max_iters
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        return [
+            base_lr
+            * 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):
+        # The new interface
+        return self.get_lr()
+
+
+class WarmupCosineLR(_LRScheduler):
+
+    def __init__(self,
+                 optimizer,
+                 max_iters,
+                 gamma=0.1,
+                 warmup_factor=0.3,
+                 warmup_iters=500,
+                 warmup_method='linear',
+                 last_epoch=-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):
+        warmup_factor = _get_warmup_factor_at_iter(
+            self.warmup_method,
+            self.last_epoch,
+            self.warmup_iters,
+            self.warmup_factor,
+        )
+        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):
+        # The new interface
+        return self.get_lr()
+
+class DummyScheduler:
+    def __init__(self, optimizer):
+        self.optimizer = optimizer
+        self.base_lrs = [group['lr'] for group in optimizer.param_groups]
+    
+    def get_last_lr(self):
+        return self.base_lrs
+
+    def step(self):
+        pass  # No operation for stepping the scheduler
+
+
+def build_scheduler(cfg, optimizer, num_iters_per_epoch):
+    """Unlike the PyTorch version, our schedulers adjust the learning rate
+    according to iteration rather than epoch.
+    """
+    # Set last epoch (here 'epoch' is actually 'iters')
+    last_epoch = (cfg.SOLVER.START_EPOCH - 1) * num_iters_per_epoch
+    
+    if cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'none' or cfg.SOLVER.SCHEDULER.SCHEDULER_NAME is None:
+        scheduler = DummyScheduler(optimizer)
+    
+    elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'multistep':
+        # Convert milestones epochs to iters
+        milestones = [(i - 1) * num_iters_per_epoch for i in cfg.SOLVER.SCHEDULER.MILESTONES]
+
+        scheduler = MultiStepLR(
+            optimizer,
+            milestones=milestones,
+            gamma=cfg.SOLVER.SCHEDULER.GAMMA,
+            last_epoch=last_epoch)
+    elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'warmup_multistep':
+        # Convert milestones epochs to iters
+        milestones = [(i - 1) * num_iters_per_epoch for i in cfg.SOLVER.SCHEDULER.MILESTONES]
+
+        # Convert warmup epochs to iters
+        cfg.SOLVER.SCHEDULER.WARMUP_ITERS = int(cfg.SOLVER.SCHEDULER.WARMUP_EPOCHS * num_iters_per_epoch)
+
+        scheduler = WarmupMultiStepLR(
+            optimizer,
+            milestones=milestones,
+            gamma=cfg.SOLVER.SCHEDULER.GAMMA,
+            warmup_factor=cfg.SOLVER.SCHEDULER.WARMUP_FACTOR,
+            warmup_iters=cfg.SOLVER.SCHEDULER.WARMUP_ITERS,
+            warmup_method=cfg.SOLVER.SCHEDULER.WARMUP_METHOD,
+            last_epoch=last_epoch)
+    elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'cosine':
+        # Get max number of iters
+        max_iters = cfg.SOLVER.NUM_EPOCHS * num_iters_per_epoch
+
+        scheduler = CosineLR(
+            optimizer,
+            max_iters=max_iters,
+            last_epoch=last_epoch)
+    elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'warmup_cosine':
+        # Get max number of iters
+        max_iters = cfg.SOLVER.NUM_EPOCHS * num_iters_per_epoch
+
+        # Convert warmup epochs to iters
+        cfg.SOLVER.SCHEDULER.WARMUP_ITERS = int(cfg.SOLVER.SCHEDULER.WARMUP_EPOCHS * num_iters_per_epoch)
+
+        scheduler = WarmupCosineLR(
+            optimizer,
+            max_iters=max_iters,
+            gamma=cfg.SOLVER.SCHEDULER.GAMMA,
+            warmup_factor=cfg.SOLVER.SCHEDULER.WARMUP_FACTOR,
+            warmup_iters=cfg.SOLVER.SCHEDULER.WARMUP_ITERS,
+            warmup_method=cfg.SOLVER.SCHEDULER.WARMUP_METHOD,
+            last_epoch=last_epoch)
+    else:
+        raise RuntimeError('Unknown lr scheduler: {}'.format(cfg.SOLVER.SCHEDULER.SCHEDULER_NAME))
+    return scheduler

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/optimizers.py

@@ -0,0 +1,27 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['build_optimizer']
+
+import torch.optim as optim
+
+
+def build_optimizer(cfg, model):
+    if cfg.SOLVER.OPTIMIZER == 'sgd':
+        optimizer = optim.SGD(
+            [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}],
+            lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY, momentum=cfg.SOLVER.MOMENTUM,
+        )
+    elif cfg.SOLVER.OPTIMIZER == 'adam':
+        optimizer = optim.Adam(
+            [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}],
+            lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+        )
+    elif cfg.SOLVER.OPTIMIZER == 'adamw':
+        optimizer = optim.AdamW(
+            [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}],
+            lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+        )
+    else:
+        raise RuntimeError('Unknown optimizer: {}'.format(cfg.SOLVER.OPTIMIZER))
+    return optimizer

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/checkpointer.py

@@ -0,0 +1,144 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['setup_checkpointer']
+
+import os.path as osp
+
+import torch
+
+
+class Checkpointer(object):
+
+    def __init__(self, cfg, phase):
+
+        # Load pretrained checkpoint
+        
+        self.is_pretrained = False
+        self.cfg = cfg
+        self.checkpoint = self._load_checkpoint(cfg.MODEL.CHECKPOINT)
+        self.except_classifier = cfg.MODEL.EXCEPT_CLASSIFIER if phase == 'train' else False
+        
+        if self.checkpoint is not None and phase == 'train':
+            # further training with the same dataset: resume
+            model_ckpt_str = cfg['MODEL']['CHECKPOINT'].replace('_', '')
+            if cfg['DATA']['DATA_NAME'].lower() in model_ckpt_str.lower():
+                cfg.SOLVER.START_EPOCH += self.checkpoint.get('epoch', 0)
+            # further training with the different dataset: finetuning on pretrained weight
+            else:
+                self.is_pretrained = True
+                print(f'WARNING) Finetuning on pretrained weight: {cfg.MODEL.CHECKPOINT}')
+        
+        elif self.checkpoint is None and phase != 'train':
+            # no checkpoint in infernece mode: error
+            raise RuntimeError('Cannot find checkpoint {}'.format(cfg.MODEL.CHECKPOINT))
+        
+        self.output_dir = cfg.OUTPUT_DIR
+
+    def load(self, model, optimizer=None):
+        if self.checkpoint is not None:
+            # try load the ckpt first, then if there is a size mismatch, do other operation
+            if not self.is_pretrained: # just resume
+                # Attempt to load the state dictionary from the checkpoint
+                model.load_state_dict(self.checkpoint['model_state_dict'])
+                # Load the optimizer state if provided
+                if optimizer is not None:
+                    if 'optimizer_state_dict' in self.checkpoint:
+                        optimizer.load_state_dict(self.checkpoint['optimizer_state_dict'])
+                    else:
+                        print('Optimizer state dictionary not found in checkpoint.')
+
+            else:
+                # different model structure (pretrained)
+                print(f'WARNING) Load checkpoint separately due to possiblity of different model structure, src/rekognition_online_action_detection/utils/checkpointer.py L51')
+                model_dict = model.state_dict()
+                unnecessary_keys = []
+                unmatched_keys = []
+                matched_ckpt = {}
+
+                # loaded_state_dict = self.checkpoint['model_state_dict']
+
+                # # Initialize a flag to track if all values match
+                # all_values_match = True
+
+                # # Iterate through all keys and compare values
+                # for key in model_dict:
+                #     import pdb; pdb.set_trace()
+                #     if not torch.allclose(model_dict[key], loaded_state_dict[key]):
+                #         print(f"Mismatch found in {key}")
+                #         print(f"Model value: {model_dict[key]}")
+                #         print(f"Loaded value: {loaded_state_dict[key]}")
+                #         print(f"Difference: {model_dict[key] - loaded_state_dict[key]}")
+                #         print("---")
+                #         all_values_match = False
+                        
+                # Filter out unnecessary keys
+                for k, v in self.checkpoint['model_state_dict'].items():
+                    if k not in model_dict:
+                        print(f'Unnecessary key: {k} not found in model')
+                        unnecessary_keys.append(k)
+                    elif model_dict[k].shape != v.shape:
+                        print(f'Unnecessary key: {k} shape {v.shape} does not match {model_dict[k].shape}')
+                        unmatched_keys.append(k)
+                    elif self.except_classifier is True and 'classifier' in k:
+                        print(f'Unnecessary key: {k} is classifier')
+                        unnecessary_keys.append(k)
+                    else:
+                        matched_ckpt[k] = v
+
+                print(f'Unloaded keys: {unnecessary_keys}')
+                print(f'Unmathced keys: {unmatched_keys}')
+                # Update the existing state dict with the matched keys
+                # import pdb; pdb.set_trace()
+                model_dict.update(matched_ckpt)
+                model.load_state_dict(model_dict)
+                # for unmatched_key in unmatched_keys:
+                #     loaded_dict_shape = self.checkpoint['model_state_dict'][unmatched_key].shape
+                #     model_dict_shape = model_dict[unmatched_key].shape
+                #     if len(loaded_dict_shape) == 2:
+                #         if loaded_dict_shape[0] != model_dict_shape[0]:
+                #             model_dict[unmatched_key][:loaded_dict_shape[0], :] = self.checkpoint['model_state_dict'][unmatched_key]
+                #             print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:{loaded_dict_shape[0]}, :] = checkpoint["model_state_dict"][{unmatched_key}]')
+                #         else:
+                #             model_dict[unmatched_key][:, :loaded_dict_shape[1]] = self.checkpoint['model_state_dict'][unmatched_key]
+                #             print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:, :{loaded_dict_shape[1]}] = checkpoint["model_state_dict"][{unmatched_key}]')
+                #     elif len(loaded_dict_shape) == 1:
+                #         model_dict[unmatched_key][:loaded_dict_shape[0]] = self.checkpoint['model_state_dict'][unmatched_key]
+                #         print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:{loaded_dict_shape[0]}] = checkpoint["model_state_dict"][{unmatched_key}]')
+                #     elif len(loaded_dict_shape) >= 3:
+                #         raise ValueError(f'Unsupported shape: {loaded_dict_shape}')
+                # model.load_state_dict(model_dict)
+                    
+                    
+
+    def save(self, epoch, model, optimizer):
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+        }, osp.join(self.output_dir, f'epoch-{epoch}.pth'))
+        print(f'\nSave checkpoint to {osp.join(self.output_dir, f"epoch-{epoch}.pth")}')
+        # osp.join(self.output_dir, f'{self.cfg.INPUT.MODALITY}_epoch-{epoch}.pth'))
+
+    def save_best(self, epoch, model, optimizer):
+        # save_path = osp.join(self.output_dir, f'{self.cfg.INPUT.MODALITY}_best.pth')
+        save_path = osp.join(self.output_dir, f'best.pth')
+        
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+        }, save_path)
+        print(f'\nSave checkpoint to {save_path}')
+
+    def _load_checkpoint(self, checkpoint):
+        if checkpoint is not None and osp.isfile(checkpoint):
+            return torch.load(checkpoint, map_location=torch.device('cpu'))
+        elif osp.isdir(checkpoint):
+            import pdb; pdb.set_trace()
+            
+        return None
+
+
+def setup_checkpointer(cfg, phase):
+    return Checkpointer(cfg, phase)

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/ek_utils.py

@@ -0,0 +1,20 @@
+import os.path as osp
+import pandas as pd
+
+
+def action_to_verb_map(path_to_data,
+                       action_offset=False,
+                       verb_offset=False):
+    actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'))
+    a_to_v = {a[1]['id'] + action_offset: a[1]['verb'] + verb_offset
+              for a in actions.iterrows()}
+    return a_to_v
+
+
+def action_to_noun_map(path_to_data,
+                       action_offset=False,
+                       noun_offset=False):
+    actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'))
+    a_to_n = {a[1]['id'] + action_offset: a[1]['noun'] + noun_offset
+              for a in actions.iterrows()}
+    return a_to_n

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/env.py

@@ -0,0 +1,32 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    'setup_random_seed',
+    'setup_environment',
+]
+
+import os
+import random
+
+import torch
+import numpy as np
+
+
+def setup_random_seed(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 setup_environment(cfg):
+    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if cfg.SEED is not None:
+        setup_random_seed(cfg.SEED)
+    return device

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/logger.py

@@ -0,0 +1,32 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ['setup_logger']
+
+import os
+import sys
+import logging
+import pprint
+
+
+def setup_logger(cfg, phase, quiet=False):
+    logger = logging.getLogger('rekognition')
+    level = logging.DEBUG if cfg.VERBOSE else logging.INFO
+    logger.setLevel(level)
+    ch = logging.StreamHandler(stream=sys.stdout)
+    ch.setLevel(level)
+    logger.addHandler(ch)
+
+    if phase == 'train':
+        log_file = os.path.join(cfg.OUTPUT_DIR, 'log.txt')
+        os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+    else:
+        log_file = os.path.splitext(cfg.MODEL.CHECKPOINT)[0] + '.txt'
+    fh = logging.FileHandler(log_file)
+    fh.setLevel(logging.INFO)
+    logger.addHandler(fh)
+
+    if not quiet:
+        logger.info(pprint.pformat(cfg))
+
+    return logger

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/parser.py

@@ -0,0 +1,114 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import json
+import os.path as osp
+
+from rekognition_online_action_detection.config.defaults import get_cfg
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Rekognition Face Anti-Spoofing')
+    parser.add_argument(
+        '--config_file',
+        default='',
+        type=str,
+        help='path to yaml config file',
+    )
+    parser.add_argument(
+        '--gpu',
+        default='0',
+        type=str,
+        help='specify visible devices'
+    )
+    parser.add_argument(
+        'opts',
+        default=None,
+        nargs='*',
+        help='modify config options using the command-line',
+    )
+    return parser.parse_args()
+
+
+def assert_and_infer_cfg(cfg, args):
+    # Setup the visible devices
+    cfg.GPU = args.gpu
+
+    # Infer data info
+    with open(cfg.DATA.DATA_INFO, 'r') as f:
+        data_info = json.load(f)[cfg.DATA.DATA_NAME]
+
+    cfg.DATA.DATA_ROOT = data_info['data_root'] if cfg.DATA.DATA_ROOT is None else cfg.DATA.DATA_ROOT
+    cfg.DATA.CLASS_NAMES = data_info['class_names'] if cfg.DATA.CLASS_NAMES is None else cfg.DATA.CLASS_NAMES
+    cfg.DATA.NUM_CLASSES = data_info['num_classes'] if cfg.DATA.NUM_CLASSES is None else cfg.DATA.NUM_CLASSES
+    cfg.DATA.IGNORE_INDEX = data_info['ignore_index'] if cfg.DATA.IGNORE_INDEX is None else cfg.DATA.IGNORE_INDEX
+    cfg.DATA.METRICS = data_info['metrics'] if cfg.DATA.METRICS is None else cfg.DATA.METRICS
+    cfg.DATA.FPS = data_info['fps'] if cfg.DATA.FPS is None else cfg.DATA.FPS
+    cfg.DATA.TRAIN_SESSION_SET = data_info['train_session_set'] if cfg.DATA.TRAIN_SESSION_SET is None else cfg.DATA.TRAIN_SESSION_SET
+    cfg.DATA.TEST_SESSION_SET = data_info['test_session_set'] if cfg.DATA.TEST_SESSION_SET is None else cfg.DATA.TEST_SESSION_SET
+
+    # Ignore two mis-labeled videos
+    if False and cfg.DATA_NAME == 'THUMOS':
+        cfg.DATA.TEST_SESSION_SET.remove('video_test_0000270')
+        cfg.DATA.TEST_SESSION_SET.remove('video_test_0001496')
+
+    # Input assertions
+    if cfg.INPUT.MODALITY == 'twostream':
+        cfg.INPUT.MODALITY = 'visual+motion'
+    if cfg.INPUT.MODALITY == 'threestream':
+        cfg.INPUT.MODALITY = 'visual+motion+object'
+    assert cfg.INPUT.MODALITY in ['visual', 'motion', 'object',
+                                  'visual+motion', 'motion+object', 'visual+object',
+                                  'visual+motion+object']
+
+    # Infer memory
+    if cfg.MODEL.MODEL_NAME in ['LSTR']:
+        cfg.MODEL.LSTR.AGES_MEMORY_LENGTH = cfg.MODEL.LSTR.AGES_MEMORY_SECONDS * cfg.DATA.FPS
+        cfg.MODEL.LSTR.LONG_MEMORY_LENGTH = cfg.MODEL.LSTR.LONG_MEMORY_SECONDS * cfg.DATA.FPS
+        cfg.MODEL.LSTR.WORK_MEMORY_LENGTH = cfg.MODEL.LSTR.WORK_MEMORY_SECONDS * cfg.DATA.FPS
+        cfg.MODEL.LSTR.TOTAL_MEMORY_LENGTH = \
+            cfg.MODEL.LSTR.AGES_MEMORY_LENGTH + \
+            cfg.MODEL.LSTR.LONG_MEMORY_LENGTH + \
+            cfg.MODEL.LSTR.WORK_MEMORY_LENGTH
+        cfg.MODEL.LSTR.ANTICIPATION_LENGTH = cfg.MODEL.LSTR.ANTICIPATION_SECONDS * cfg.DATA.FPS
+        assert cfg.MODEL.LSTR.AGES_MEMORY_LENGTH % cfg.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE == 0
+        assert cfg.MODEL.LSTR.LONG_MEMORY_LENGTH % cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE == 0
+        assert cfg.MODEL.LSTR.WORK_MEMORY_LENGTH % cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE == 0
+        assert cfg.MODEL.LSTR.ANTICIPATION_LENGTH % cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE == 0
+        cfg.MODEL.LSTR.AGES_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.AGES_MEMORY_LENGTH // cfg.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE
+        cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH // cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE
+        cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH // cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE
+        cfg.MODEL.LSTR.TOTAL_MEMORY_NUM_SAMPLES = \
+            cfg.MODEL.LSTR.AGES_MEMORY_NUM_SAMPLES + \
+            cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + \
+            cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES
+        cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES = cfg.MODEL.LSTR.ANTICIPATION_LENGTH // cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE
+
+        assert cfg.MODEL.LSTR.INFERENCE_MODE in ['batch', 'stream']
+
+    # Infer output dir
+    config_name = osp.splitext(args.config_file)[0].split('/')[1:]
+    cfg.OUTPUT_DIR = osp.join(cfg.OUTPUT_DIR, *config_name)
+    if cfg.SESSION:
+        cfg.OUTPUT_DIR = osp.join(cfg.OUTPUT_DIR, cfg.SESSION)
+
+
+def load_cfg():
+    args = parse_args()
+    cfg = get_cfg()
+    
+    for i in range(0, len(args.opts), 2):
+        key = args.opts[i]
+        value = args.opts[i+1]
+        if key == 'MODEL.LSTR.LONG_MEMORY_SECONDS' or key == 'MODEL.LSTR.WORK_MEMORY_SECONDS':
+            args.opts[i+1] = int(value)
+            print(f'(key, value): ({key}, {value}) is converted to integer type')
+    
+    if args.config_file is not None:
+        cfg.merge_from_file(args.config_file)
+    if args.opts is not None:
+        cfg.merge_from_list(args.opts)
+    assert_and_infer_cfg(cfg, args)
+    # config_file is in cfg.OUTPUT_DIR or args.config_file
+    return cfg

third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/registry.py

@@ -0,0 +1,23 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+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