A fast and high-performance visual object tracker with real-time speed on Jetson.
from easydict import EasyDict as edict
import yaml

"""
Add default config for OSTrack.
"""
cfg = edict()

# MODEL
cfg.MODEL = edict()
cfg.MODEL.PRETRAIN_FILE = "cae_base.pth"
cfg.MODEL.AUX_TEMPLATE = False
cfg.MODEL.USE_CLS_TOKEN = True
cfg.MODEL.USE_QUALITY_SCORE = False
cfg.MODEL.ADD_TEMPLATE_TARGET_EMBED = False

# MODEL.BACKBONE
cfg.MODEL.BACKBONE = edict()
cfg.MODEL.BACKBONE.TYPE = "vit_base_patch16_224_CAE"
cfg.MODEL.BACKBONE.STRIDE = 16
cfg.MODEL.BACKBONE.ADD_CLS_TOKEN = False
cfg.MODEL.BACKBONE.CLS_TOKEN_USE_MODE = 'ignore'
cfg.MODEL.BACKBONE.AI_LAYERS = 3
cfg.MODEL.BACKBONE.DEPTH = 9

cfg.MODEL.BACKBONE.CE = False
cfg.MODEL.BACKBONE.CE_LOC = []

# MODEL.HEAD
cfg.MODEL.HEAD = edict()
cfg.MODEL.HEAD.TYPE = "CENTER"
cfg.MODEL.HEAD.NUM_CHANNELS = 256


# TRAIN
cfg.TRAIN = edict()
cfg.TRAIN.LR = 0.0001
cfg.TRAIN.WEIGHT_DECAY = 0.0001
cfg.TRAIN.EPOCH = 100
cfg.TRAIN.LR_DROP_EPOCH = 80
cfg.TRAIN.BATCH_SIZE = 32
cfg.TRAIN.NUM_WORKER = 8
cfg.TRAIN.OPTIMIZER = "ADAMW"
cfg.TRAIN.BACKBONE_MULTIPLIER = 0.1
cfg.TRAIN.GIOU_WEIGHT = 2.0
cfg.TRAIN.L1_WEIGHT = 5.0
cfg.TRAIN.FREEZE_LAYERS = [0, ]
cfg.TRAIN.PRINT_INTERVAL = 50
cfg.TRAIN.VAL_EPOCH_INTERVAL = 20
cfg.TRAIN.GRAD_CLIP_NORM = 0.1
cfg.TRAIN.AMP = False
cfg.TRAIN.TRAIN_CLS = False
cfg.TRAIN.POSITIVE_PROB = 0.5
cfg.TRAIN.DROP_PATH_RATE = 0.1  # drop path rate for ViT backbone

# TRAIN.SCHEDULER
cfg.TRAIN.SCHEDULER = edict()
cfg.TRAIN.SCHEDULER.TYPE = "step"
cfg.TRAIN.SCHEDULER.DECAY_RATE = 0.1

# DATA
cfg.DATA = edict()
cfg.DATA.SAMPLER_MODE = "causal"  # sampling methods
cfg.DATA.MEAN = [0.485, 0.456, 0.406]
cfg.DATA.STD = [0.229, 0.224, 0.225]
cfg.DATA.MAX_SAMPLE_INTERVAL = 200
# DATA.TRAIN
cfg.DATA.TRAIN = edict()
cfg.DATA.TRAIN.DATASETS_NAME = ["LASOT", "GOT10K_vottrain"]
cfg.DATA.TRAIN.DATASETS_RATIO = [1, 1]
cfg.DATA.TRAIN.SAMPLE_PER_EPOCH = 60000
# DATA.VAL
cfg.DATA.VAL = edict()
cfg.DATA.VAL.DATASETS_NAME = ["GOT10K_votval"]
cfg.DATA.VAL.DATASETS_RATIO = [1]
cfg.DATA.VAL.SAMPLE_PER_EPOCH = 10000
# DATA.SEARCH
cfg.DATA.SEARCH = edict()
cfg.DATA.SEARCH.SIZE = 320
cfg.DATA.SEARCH.FACTOR = 5.0
cfg.DATA.SEARCH.CENTER_JITTER = 4.5
cfg.DATA.SEARCH.SCALE_JITTER = 0.5
cfg.DATA.SEARCH.NUMBER = 1
# DATA.TEMPLATE
cfg.DATA.TEMPLATE = edict()
cfg.DATA.TEMPLATE.NUMBER = 1
cfg.DATA.TEMPLATE.SIZE = 128
cfg.DATA.TEMPLATE.FACTOR = 2.0
cfg.DATA.TEMPLATE.CENTER_JITTER = 0
cfg.DATA.TEMPLATE.SCALE_JITTER = 0

# TEST
cfg.TEST = edict()
cfg.TEST.TEMPLATE_FACTOR = 2.0
cfg.TEST.TEMPLATE_SIZE = 128
cfg.TEST.SEARCH_FACTOR = 5.0
cfg.TEST.SEARCH_SIZE = 320
cfg.TEST.EPOCH = 500


def _edict2dict(dest_dict, src_edict):
    if isinstance(dest_dict, dict) and isinstance(src_edict, dict):
        for k, v in src_edict.items():
            if not isinstance(v, edict):
                dest_dict[k] = v
            else:
                dest_dict[k] = {}
                _edict2dict(dest_dict[k], v)
    else:
        return


def gen_config(config_file):
    cfg_dict = {}
    _edict2dict(cfg_dict, cfg)
    with open(config_file, 'w') as f:
        yaml.dump(cfg_dict, f, default_flow_style=False)


def _update_config(base_cfg, exp_cfg):
    if isinstance(base_cfg, dict) and isinstance(exp_cfg, edict):
        for k, v in exp_cfg.items():
            if k in base_cfg:
                if not isinstance(v, dict):
                    base_cfg[k] = v
                else:
                    _update_config(base_cfg[k], v)
            else:
                raise ValueError("{} not exist in config.py".format(k))
    else:
        return


def update_config_from_file(filename, base_cfg=None):
    exp_config = None
    with open(filename) as f:
        exp_config = edict(yaml.safe_load(f))
        if base_cfg is not None:
            _update_config(base_cfg, exp_config)
        else:
            _update_config(cfg, exp_config)

lib/config/litetrack/config.py
TsingWei/LiteTrack
from easydict import EasyDict as edict
import yaml

"""
Add default config for OSTrack.
"""
cfg = edict()

# MODEL
cfg.MODEL = edict()
cfg.MODEL.PRETRAIN_FILE = "mae_pretrain_vit_base.pth"
cfg.MODEL.EXTRA_MERGER = False
cfg.MODEL.AUX_TEMPLATE = False

cfg.MODEL.RETURN_INTER = False
cfg.MODEL.RETURN_STAGES = [2, 5, 8, 11]
cfg.MODEL.USE_CLS_TOKEN = False
cfg.MODEL.USE_QUALITY_SCORE = False
cfg.MODEL.ADD_TEMPLATE_TARGET_EMBED = False

# MODEL.BACKBONE
cfg.MODEL.BACKBONE = edict()
cfg.MODEL.BACKBONE.TYPE = "vit_base_patch16_224"
cfg.MODEL.BACKBONE.STRIDE = 16
cfg.MODEL.BACKBONE.MID_PE = False
cfg.MODEL.BACKBONE.SEP_SEG = False
cfg.MODEL.BACKBONE.CAT_MODE = 'direct'
cfg.MODEL.BACKBONE.MERGE_LAYER = 0
cfg.MODEL.BACKBONE.ADD_CLS_TOKEN = False
cfg.MODEL.BACKBONE.CLS_TOKEN_USE_MODE = 'ignore'

cfg.MODEL.BACKBONE.CE = False
cfg.MODEL.BACKBONE.CE_LOC = []
cfg.MODEL.BACKBONE.CE_KEEP_RATIO = []
cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE = 'ALL'  # choose between ALL, CTR_POINT, CTR_REC, GT_BOX
cfg.MODEL.BACKBONE.SEP_LOC = []
# MODEL.HEAD
cfg.MODEL.HEAD = edict()
cfg.MODEL.HEAD.TYPE = "CENTER"
cfg.MODEL.HEAD.NUM_CHANNELS = 768
cfg.MODEL.HEAD.NUM_QUERIES = 30
cfg.MODEL.HEAD.NUM_LAYERS = 2
cfg.MODEL.HEAD.ATTN = 'SELF'
cfg.MODEL.HEAD.CLASSIFICATION_POSITION = 'FFN'

# TRAIN
cfg.TRAIN = edict()
cfg.TRAIN.LR = 0.0001
cfg.TRAIN.WEIGHT_DECAY = 0.0001
cfg.TRAIN.EPOCH = 500
cfg.TRAIN.LR_DROP_EPOCH = 400
cfg.TRAIN.BATCH_SIZE = 16
cfg.TRAIN.NUM_WORKER = 8
cfg.TRAIN.OPTIMIZER = "ADAMW"
cfg.TRAIN.BACKBONE_MULTIPLIER = 0.1
cfg.TRAIN.GIOU_WEIGHT = 2.0
cfg.TRAIN.L1_WEIGHT = 5.0
cfg.TRAIN.FREEZE_LAYERS = [0, ]
cfg.TRAIN.PRINT_INTERVAL = 50
cfg.TRAIN.VAL_EPOCH_INTERVAL = 20
cfg.TRAIN.GRAD_CLIP_NORM = 0.1
cfg.TRAIN.AMP = False
cfg.TRAIN.TRAIN_CLS = False
cfg.TRAIN.POSITIVE_PROB = 0.5

cfg.TRAIN.CE_START_EPOCH = 20  # candidate elimination start epoch
cfg.TRAIN.CE_WARM_EPOCH = 80  # candidate elimination warm up epoch
cfg.TRAIN.DROP_PATH_RATE = 0.1  # drop path rate for ViT backbone

# TRAIN.SCHEDULER
cfg.TRAIN.SCHEDULER = edict()
cfg.TRAIN.SCHEDULER.TYPE = "step"
cfg.TRAIN.SCHEDULER.DECAY_RATE = 0.1

# DATA
cfg.DATA = edict()
cfg.DATA.SAMPLER_MODE = "causal"  # sampling methods
cfg.DATA.MEAN = [0.485, 0.456, 0.406]
cfg.DATA.STD = [0.229, 0.224, 0.225]
cfg.DATA.MAX_SAMPLE_INTERVAL = 200
# DATA.TRAIN
cfg.DATA.TRAIN = edict()
cfg.DATA.TRAIN.DATASETS_NAME = ["LASOT", "GOT10K_vottrain"]
cfg.DATA.TRAIN.DATASETS_RATIO = [1, 1]
cfg.DATA.TRAIN.SAMPLE_PER_EPOCH = 60000
# DATA.VAL
cfg.DATA.VAL = edict()
cfg.DATA.VAL.DATASETS_NAME = ["GOT10K_votval"]
cfg.DATA.VAL.DATASETS_RATIO = [1]
cfg.DATA.VAL.SAMPLE_PER_EPOCH = 10000
# DATA.SEARCH
cfg.DATA.SEARCH = edict()
cfg.DATA.SEARCH.SIZE = 320
cfg.DATA.SEARCH.FACTOR = 5.0
cfg.DATA.SEARCH.CENTER_JITTER = 4.5
cfg.DATA.SEARCH.SCALE_JITTER = 0.5
cfg.DATA.SEARCH.NUMBER = 1
# DATA.TEMPLATE
cfg.DATA.TEMPLATE = edict()
cfg.DATA.TEMPLATE.NUMBER = 1
cfg.DATA.TEMPLATE.SIZE = 128
cfg.DATA.TEMPLATE.FACTOR = 2.0
cfg.DATA.TEMPLATE.CENTER_JITTER = 0
cfg.DATA.TEMPLATE.SCALE_JITTER = 0

# TEST
cfg.TEST = edict()
cfg.TEST.TEMPLATE_FACTOR = 2.0
cfg.TEST.TEMPLATE_SIZE = 128
cfg.TEST.SEARCH_FACTOR = 5.0
cfg.TEST.SEARCH_SIZE = 320
cfg.TEST.EPOCH = 500


def _edict2dict(dest_dict, src_edict):
    if isinstance(dest_dict, dict) and isinstance(src_edict, dict):
        for k, v in src_edict.items():
            if not isinstance(v, edict):
                dest_dict[k] = v
            else:
                dest_dict[k] = {}
                _edict2dict(dest_dict[k], v)
    else:
        return


def gen_config(config_file):
    cfg_dict = {}
    _edict2dict(cfg_dict, cfg)
    with open(config_file, 'w') as f:
        yaml.dump(cfg_dict, f, default_flow_style=False)


def _update_config(base_cfg, exp_cfg):
    if isinstance(base_cfg, dict) and isinstance(exp_cfg, edict):
        for k, v in exp_cfg.items():
            if k in base_cfg:
                if not isinstance(v, dict):
                    base_cfg[k] = v
                else:
                    _update_config(base_cfg[k], v)
            else:
                raise ValueError("{} not exist in config.py".format(k))
    else:
        return


def update_config_from_file(filename, base_cfg=None):
    exp_config = None
    with open(filename) as f:
        exp_config = edict(yaml.safe_load(f))
        if base_cfg is not None:
            _update_config(base_cfg, exp_config)
        else:
            _update_config(cfg, exp_config)

lib/config/ostrack/config.py
TsingWei/LiteTrack
from .ostrack.ostrack import build_ostrack
from .litetrack.litetrack import build_LiteTrack
lib/models/__init__.py
TsingWei/LiteTrack
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_

from lib.models.layers.rpe import generate_2d_concatenated_self_attention_relative_positional_encoding_index


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.,
                 rpe=False, z_size=8, x_size=16):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.rpe =rpe
        if self.rpe:
            relative_position_index = \
                generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
                                                                                           [x_size, x_size])
            self.register_buffer("relative_position_index", relative_position_index)
            # define a parameter table of relative position bias
            self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
                                                                          relative_position_index.max() + 1)))
            trunc_normal_(self.relative_position_bias_table, std=0.02)

    def forward(self, x, mask=None, return_attention=False):
        # x: B, N, C
        # mask: [B, N, ] torch.bool
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = self.attn_drop(attn)

        x = (attn @ v)
        # x = F.scaled_dot_product_attention(q, k, v)
        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        if return_attention:
            return x, attn
        else:
            return x


class Attention_talking_head(nn.Module):
    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
    # with slight modifications to add Talking Heads Attention (https://arxiv.org/pdf/2003.02436v1.pdf)
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                 rpe=True, z_size=7, x_size=14):
        super().__init__()

        self.num_heads = num_heads

        head_dim = dim // num_heads

        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)

        self.proj = nn.Linear(dim, dim)

        self.proj_l = nn.Linear(num_heads, num_heads)
        self.proj_w = nn.Linear(num_heads, num_heads)

        self.proj_drop = nn.Dropout(proj_drop)

        self.rpe = rpe
        if self.rpe:
            relative_position_index = \
                generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
                                                                                           [x_size, x_size])
            self.register_buffer("relative_position_index", relative_position_index)
            # define a parameter table of relative position bias
            self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
                                                                          relative_position_index.max() + 1)))
            trunc_normal_(self.relative_position_bias_table, std=0.02)

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1))

        if self.rpe:
            relative_position_bias = self.relative_position_bias_table[:, self.relative_position_index].unsqueeze(0)
            attn += relative_position_bias

        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2),
                                    float('-inf'),)

        attn = self.proj_l(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

        attn = attn.softmax(dim=-1)

        attn = self.proj_w(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
lib/models/layers/attn.py
TsingWei/LiteTrack
import math
import torch
import torch.nn as nn
from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_

from lib.models.layers.attn import Attention


def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor, lens_t: int, keep_ratio: float, global_index: torch.Tensor, box_mask_z: torch.Tensor):
    """
    Eliminate potential background candidates for computation reduction and noise cancellation.
    Args:
        attn (torch.Tensor): [B, num_heads, L_t + L_s, L_t + L_s], attention weights
        tokens (torch.Tensor):  [B, L_t + L_s, C], template and search region tokens
        lens_t (int): length of template
        keep_ratio (float): keep ratio of search region tokens (candidates)
        global_index (torch.Tensor): global index of search region tokens
        box_mask_z (torch.Tensor): template mask used to accumulate attention weights

    Returns:
        tokens_new (torch.Tensor): tokens after candidate elimination
        keep_index (torch.Tensor): indices of kept search region tokens
        removed_index (torch.Tensor): indices of removed search region tokens
    """
    lens_s = attn.shape[-1] - lens_t
    bs, hn, _, _ = attn.shape

    lens_keep = math.ceil(keep_ratio * lens_s)
    if lens_keep == lens_s:
        return tokens, global_index, None

    attn_t = attn[:, :, :lens_t, lens_t:]

    if box_mask_z is not None:
        box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(-1, attn_t.shape[1], -1, attn_t.shape[-1])
        # attn_t = attn_t[:, :, box_mask_z, :]
        attn_t = attn_t[box_mask_z]
        attn_t = attn_t.view(bs, hn, -1, lens_s)
        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s

        # attn_t = [attn_t[i, :, box_mask_z[i, :], :] for i in range(attn_t.size(0))]
        # attn_t = [attn_t[i].mean(dim=1).mean(dim=0) for i in range(len(attn_t))]
        # attn_t = torch.stack(attn_t, dim=0)
    else:
        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s

    # use sort instead of topk, due to the speed issue
    # https://github.com/pytorch/pytorch/issues/22812
    sorted_attn, indices = torch.sort(attn_t, dim=1, descending=True)

    topk_attn, topk_idx = sorted_attn[:, :lens_keep], indices[:, :lens_keep]
    non_topk_attn, non_topk_idx = sorted_attn[:, lens_keep:], indices[:, lens_keep:]

    keep_index = global_index.gather(dim=1, index=topk_idx)
    removed_index = global_index.gather(dim=1, index=non_topk_idx)

    # separate template and search tokens
    tokens_t = tokens[:, :lens_t]
    tokens_s = tokens[:, lens_t:]

    # obtain the attentive and inattentive tokens
    B, L, C = tokens_s.shape
    # topk_idx_ = topk_idx.unsqueeze(-1).expand(B, lens_keep, C)
    attentive_tokens = tokens_s.gather(dim=1, index=topk_idx.unsqueeze(-1).expand(B, -1, C))
    # inattentive_tokens = tokens_s.gather(dim=1, index=non_topk_idx.unsqueeze(-1).expand(B, -1, C))

    # compute the weighted combination of inattentive tokens
    # fused_token = non_topk_attn @ inattentive_tokens

    # concatenate these tokens
    # tokens_new = torch.cat([tokens_t, attentive_tokens, fused_token], dim=0)
    tokens_new = torch.cat([tokens_t, attentive_tokens], dim=1)

    return tokens_new, keep_index, removed_index


class CEBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, keep_ratio_search=1.0,):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

        self.keep_ratio_search = keep_ratio_search

    def forward(self, x, global_index_template, global_index_search, mask=None, ce_template_mask=None, keep_ratio_search=None):
        x_attn, attn = self.attn(self.norm1(x), mask, True)
        x = x + self.drop_path(x_attn)
        lens_t = global_index_template.shape[1]

        removed_index_search = None
        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)

        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x, global_index_template, global_index_search, removed_index_search, attn


class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, mask=None):
        x = x + self.drop_path(self.attn(self.norm1(x), mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

lib/models/layers/attn_blocks.py
TsingWei/LiteTrack
import torch.nn as nn
import torch
import torch.nn.functional as F
from einops import rearrange
import numpy as np
# from mmengine.model import constant_init, normal_init
# from mmcv.cnn import ConvModule
# from mmcv.ops import DeformConv2d
# from .decoder_head import DecoderHead
# from mmcv.ops import 
# from mmdet.core.utils import DeformConv

from lib.models.layers.frozen_bn import FrozenBatchNorm2d

class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """
    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError 
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x
        
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1,
         freeze_bn=False):
    if freeze_bn:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True),
            FrozenBatchNorm2d(out_planes),
            nn.ReLU(inplace=True))
    else:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True),
            nn.BatchNorm2d(out_planes),
            nn.ReLU(inplace=True))

def c_sigmoid(x):
    y = torch.clamp(x.sigmoid_(), min=1e-4, max=1-1e-4)
    return y

class Corner_Predictor(nn.Module):
    """ Corner Predictor module"""

    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
        super(Corner_Predictor, self).__init__()
        self.feat_sz = feat_sz
        self.stride = stride
        self.img_sz = self.feat_sz * self.stride
        self.grid_size = self.feat_sz+1
        '''top-left corner'''
        self.conv1_tl = conv(inplanes, channel, freeze_bn=freeze_bn, kernel_size=2)
        self.conv2_tl = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_tl = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_tl = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_tl = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''bottom-right corner'''
        self.conv1_br = conv(inplanes, channel, freeze_bn=freeze_bn, kernel_size=2)
        self.conv2_br = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_br = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_br = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_br = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''about coordinates and indexs'''
        with torch.no_grad():
            
            self.indice = torch.arange(0, self.grid_size ).view(-1, 1) * self.stride
            # generate mesh-grid
            self.coord_x = self.indice.repeat((self.grid_size , 1)) \
                .view((self.grid_size  * self.grid_size ,)).float().cuda()
            self.coord_y = self.indice.repeat((1, self.grid_size )) \
                .view((self.grid_size  * self.grid_size ,)).float().cuda()

    def forward(self, x, return_dist=False, softmax=True):
        """ Forward pass with input x. """
        score_map_tl, score_map_br = self.get_score_map(x)
        if return_dist:
            coorx_tl, coory_tl, prob_vec_tl = self.soft_argmax(score_map_tl, return_dist=True, softmax=softmax)
            coorx_br, coory_br, prob_vec_br = self.soft_argmax(score_map_br, return_dist=True, softmax=softmax)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz, torch.stack((prob_vec_tl,prob_vec_br),1)
        else:
            coorx_tl, coory_tl = self.soft_argmax(score_map_tl)
            coorx_br, coory_br = self.soft_argmax(score_map_br)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz

    def get_score_map(self, x):
        # top-left branch
        x_tl1 = self.conv1_tl(x)
        x_tl2 = self.conv2_tl(x_tl1)
        x_tl3 = self.conv3_tl(x_tl2)
        x_tl4 = self.conv4_tl(x_tl3)
        score_map_tl = self.conv5_tl(x_tl4)

        # bottom-right branch
        x_br1 = self.conv1_br(x)
        x_br2 = self.conv2_br(x_br1)
        x_br3 = self.conv3_br(x_br2)
        x_br4 = self.conv4_br(x_br3)
        score_map_br = self.conv5_br(x_br4)
        return score_map_tl, score_map_br

    def soft_argmax(self, score_map, return_dist=False, softmax=True):
        """ get soft-argmax coordinate for a given heatmap """
        score_vec = score_map.view((-1, self.grid_size  * self.grid_size ))  # (batch, feat_sz * feat_sz)
        prob_vec = nn.functional.softmax(score_vec, dim=1)
        exp_x = torch.sum((self.coord_x * prob_vec), dim=1)
        exp_y = torch.sum((self.coord_y * prob_vec), dim=1)
        if return_dist:
            if softmax:
                return exp_x, exp_y, prob_vec
            else:
                return exp_x, exp_y, score_vec
        else:
            return exp_x, exp_y


class CenterPredictor(nn.Module, ):
    def __init__(self, inplanes=64, channel=256, feat_sz=20, feat_tz=20, stride=16, freeze_bn=False):
        super(CenterPredictor, self).__init__()
        # self.feat_sz = feat_sz
        # self.feat_tz = feat_tz
        self.stride = stride
        # self.img_sz = self.feat_sz * self.stride

        # corner predict
        self.conv1_ctr = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_ctr = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_ctr = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_ctr = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_ctr = nn.Conv2d(channel // 8, 1, kernel_size=1)

        # size regress
        self.conv1_offset = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_offset = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_offset = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_offset = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_offset = nn.Conv2d(channel // 8, 2, kernel_size=1)

        # size regress
        self.conv1_size = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_size = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_size = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_size = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_size = nn.Conv2d(channel // 8, 2, kernel_size=1)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, x, feat_size, gt_score_map=None):
        """ Forward pass with input x. """
        score_map_ctr, size_map, offset_map = self.get_score_map(x)

        # assert gt_score_map is None
        if True:
            bbox = self.cal_bbox(score_map_ctr, feat_size, size_map, offset_map)
        else:
            bbox = self.cal_bbox(gt_score_map.unsqueeze(1), size_map, offset_map)

        return score_map_ctr, bbox, size_map, offset_map

    def cal_bbox(self, score_map_ctr, feat_size, size_map, offset_map, return_score=False):
        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
        idx_y = idx // feat_size
        idx_x = idx % feat_size

        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
        size = size_map.flatten(2).gather(dim=2, index=idx)
        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)

        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
        # cx, cy, w, h
        bbox = torch.cat([(idx_x.to(torch.float) + offset[:, :1]) / feat_size,
                          (idx_y.to(torch.float) + offset[:, 1:]) / feat_size,
                          size.squeeze(-1)], dim=1)

        if return_score:
            return bbox, max_score
        return bbox

    def get_pred(self, score_map_ctr, size_map, offset_map):
        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
        idx_y = idx // self.feat_sz
        idx_x = idx % self.feat_sz

        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
        size = size_map.flatten(2).gather(dim=2, index=idx)
        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)

        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
        return size * self.feat_sz, offset

    def get_score_map(self, x):

        def _sigmoid(x):
            y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4)
            return y

        # ctr branch
        x_ctr1 = self.conv1_ctr(x)
        x_ctr2 = self.conv2_ctr(x_ctr1)
        x_ctr3 = self.conv3_ctr(x_ctr2)
        x_ctr4 = self.conv4_ctr(x_ctr3)
        score_map_ctr = self.conv5_ctr(x_ctr4)

        # offset branch
        x_offset1 = self.conv1_offset(x)
        x_offset2 = self.conv2_offset(x_offset1)
        x_offset3 = self.conv3_offset(x_offset2)
        x_offset4 = self.conv4_offset(x_offset3)
        score_map_offset = self.conv5_offset(x_offset4)

        # size branch
        x_size1 = self.conv1_size(x)
        x_size2 = self.conv2_size(x_size1)
        x_size3 = self.conv3_size(x_size2)
        x_size4 = self.conv4_size(x_size3)
        score_map_size = self.conv5_size(x_size4)
        return _sigmoid(score_map_ctr), _sigmoid(score_map_size), score_map_offset

def build_box_head(cfg, hidden_dim):
    stride = cfg.MODEL.BACKBONE.STRIDE

    if "CORNER" in cfg.MODEL.HEAD.TYPE:
        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
        channel = getattr(cfg.MODEL, "NUM_CHANNELS", 256)
        print("head channel: %d" % channel)
        if cfg.MODEL.HEAD.TYPE == "CORNER":
            corner_head = Corner_Predictor(inplanes=hidden_dim, channel=channel,
                                           feat_sz=feat_sz, stride=stride)
        else:
            raise ValueError()
        return corner_head
    elif cfg.MODEL.HEAD.TYPE == "CENTER":
        in_channel = hidden_dim
        out_channel = cfg.MODEL.HEAD.NUM_CHANNELS
        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
        center_head = CenterPredictor(inplanes=in_channel, channel=out_channel,
                                      feat_sz=feat_sz, stride=stride)
        return center_head
    else:
        raise ValueError("HEAD TYPE %s is not supported." % cfg.MODEL.HEAD_TYPE)

lib/models/layers/head.py
TsingWei/LiteTrack
import torch.nn as nn

from timm.models.layers import to_2tuple


class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
        # allow different input size
        # B, C, H, W = x.shape
        # _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
        # _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x

lib/models/layers/patch_embed.py
TsingWei/LiteTrack
import torch
import torch.nn as nn
from timm.models.layers import trunc_normal_


def generate_2d_relative_positional_encoding_index(z_shape, x_shape):
    '''
        z_shape: (z_h, z_w)
        x_shape: (x_h, x_w)
    '''
    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))

    z_2d_index_h = z_2d_index_h.flatten(0)
    z_2d_index_w = z_2d_index_w.flatten(0)
    x_2d_index_h = x_2d_index_h.flatten(0)
    x_2d_index_w = x_2d_index_w.flatten(0)

    diff_h = z_2d_index_h[:, None] - x_2d_index_h[None, :]
    diff_w = z_2d_index_w[:, None] - x_2d_index_w[None, :]

    diff = torch.stack((diff_h, diff_w), dim=-1)
    _, indices = torch.unique(diff.view(-1, 2), return_inverse=True, dim=0)
    return indices.view(z_shape[0] * z_shape[1], x_shape[0] * x_shape[1])


def generate_2d_concatenated_self_attention_relative_positional_encoding_index(z_shape, x_shape):
    '''
        z_shape: (z_h, z_w)
        x_shape: (x_h, x_w)
    '''
    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))

    z_2d_index_h = z_2d_index_h.flatten(0)
    z_2d_index_w = z_2d_index_w.flatten(0)
    x_2d_index_h = x_2d_index_h.flatten(0)
    x_2d_index_w = x_2d_index_w.flatten(0)

    concatenated_2d_index_h = torch.cat((z_2d_index_h, x_2d_index_h))
    concatenated_2d_index_w = torch.cat((z_2d_index_w, x_2d_index_w))

    diff_h = concatenated_2d_index_h[:, None] - concatenated_2d_index_h[None, :]
    diff_w = concatenated_2d_index_w[:, None] - concatenated_2d_index_w[None, :]

    z_len = z_shape[0] * z_shape[1]
    x_len = x_shape[0] * x_shape[1]
    a = torch.empty((z_len + x_len), dtype=torch.int64)
    a[:z_len] = 0
    a[z_len:] = 1
    b=a[:, None].repeat(1, z_len + x_len)
    c=a[None, :].repeat(z_len + x_len, 1)

    diff = torch.stack((diff_h, diff_w, b, c), dim=-1)
    _, indices = torch.unique(diff.view((z_len + x_len) * (z_len + x_len), 4), return_inverse=True, dim=0)
    return indices.view((z_len + x_len), (z_len + x_len))


def generate_2d_concatenated_cross_attention_relative_positional_encoding_index(z_shape, x_shape):
    '''
        z_shape: (z_h, z_w)
        x_shape: (x_h, x_w)
    '''
    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))

    z_2d_index_h = z_2d_index_h.flatten(0)
    z_2d_index_w = z_2d_index_w.flatten(0)
    x_2d_index_h = x_2d_index_h.flatten(0)
    x_2d_index_w = x_2d_index_w.flatten(0)

    concatenated_2d_index_h = torch.cat((z_2d_index_h, x_2d_index_h))
    concatenated_2d_index_w = torch.cat((z_2d_index_w, x_2d_index_w))

    diff_h = x_2d_index_h[:, None] - concatenated_2d_index_h[None, :]
    diff_w = x_2d_index_w[:, None] - concatenated_2d_index_w[None, :]

    z_len = z_shape[0] * z_shape[1]
    x_len = x_shape[0] * x_shape[1]

    a = torch.empty(z_len + x_len, dtype=torch.int64)
    a[: z_len] = 0
    a[z_len:] = 1
    c = a[None, :].repeat(x_len, 1)

    diff = torch.stack((diff_h, diff_w, c), dim=-1)
    _, indices = torch.unique(diff.view(x_len * (z_len + x_len), 3), return_inverse=True, dim=0)
    return indices.view(x_len, (z_len + x_len))


class RelativePosition2DEncoder(nn.Module):
    def __init__(self, num_heads, embed_size):
        super(RelativePosition2DEncoder, self).__init__()
        self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads, embed_size)))
        trunc_normal_(self.relative_position_bias_table, std=0.02)

    def forward(self, attn_rpe_index):
        '''
            Args:
                attn_rpe_index (torch.Tensor): (*), any shape containing indices, max(attn_rpe_index) < embed_size
            Returns:
                torch.Tensor: (1, num_heads, *)
        '''
        return self.relative_position_bias_table[:, attn_rpe_index].unsqueeze(0)

lib/models/layers/rpe.py
TsingWei/LiteTrack
"""
Basic LiteTrack model.
"""
import math
import os
from typing import List

import torch
from torch import nn
from torch.nn.modules.transformer import _get_clones

from lib.models.layers.head import build_box_head
from lib.models.litetrack.vit_cae_async import CAE_Base_patch16_224_Async
from lib.utils.box_ops import box_xyxy_to_cxcywh
from lib.test.utils.hann import hann2d


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class LiteTrack(nn.Module):
    """ This is the base class for LiteTrack """

    def __init__(self, transformer, box_head, aux_loss=False, head_type="CORNER", aux_template_loss=False,
                 search_feat_size=16, template_feat_size=8, add_target_token=False):
        """ Initializes the model.
        Parameters:
            transformer: torch module of the transformer architecture.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.backbone = transformer
        self.box_head = box_head

        self.aux_loss = aux_loss
        self.head_type = head_type

        self.feat_size_s = int(search_feat_size)
        self.feat_len_s = int(search_feat_size ** 2)
        self.feat_size_t = int(template_feat_size)
        self.feat_len_t = int(template_feat_size ** 2)

        if self.aux_loss:
            self.box_head = _get_clones(self.box_head, 6)
        self.aux_template_loss = aux_template_loss
        self.add_target_token = add_target_token
        if self.add_target_token:
            self.target_token_embed = Mlp(
                4, out_features=self.backbone.embed_dim)
        self.onnx = False
        if self.onnx:
            self.output_window = hann2d(torch.tensor([self.feat_sz, self.feat_sz]).long(), centered=True)

    def forward_z(self, template: torch.Tensor, template_bb=None,):
        target_token = None
        if self.add_target_token:
            target_token = self.target_token_embed(template_bb).unsqueeze(-2)
        return self.backbone(z=template, target_token=target_token, mode='z')

    def forward(self,template_feats, search ):
        x = self.backbone(z=template_feats,x=search, mode='x')
        out = self.forward_head(x)
        if self.onnx:
            response = self.output_window * out['score_map']
            return response, out['size_map'], out['offset_map']
        return out

    def forward_train(self, template: torch.Tensor,
                search: torch.Tensor,
                template_bb=None,
                gt_cls_map=None,
                ):
        if self.add_target_token:
            target_token = self.target_token_embed(template_bb).unsqueeze(-2)
            x= self.backbone(z=template,
                                        x=search,
                                        target_token=target_token
                                        )
        else:
            x = self.backbone(z=template, x=search)
        out = self.forward_head(x, gt_cls_map)

        return out

    def forward_head(self, cat_feature,
                     gt_score_map=None,
                     ):
        """
        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
        """
        if self.head_type == "CENTER":
            # encoder output for the search region (B, HW, C)
            enc_opt = cat_feature[:, -self.feat_len_s:]
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)

            score_map_ctr, bbox, size_map, offset_map = self.box_head(
                opt_feat, self.feat_size_s)
            # outputs_coord = box_xyxy_to_cxcywh(bbox)
            outputs_coord = bbox
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map_ctr,
                   'size_map': size_map,
                   'offset_map': offset_map}

            return out

        else:
            raise NotImplementedError


def build_LiteTrack(cfg, training=True):
    current_dir = os.path.dirname(os.path.abspath(
        __file__))  # This is your Project Root
    pretrained_path = os.path.join(current_dir, '../../../pretrained_models')
    if cfg.MODEL.PRETRAIN_FILE and ('LiteTrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
    else:
        pretrained = ''

    if cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224_CAE':
        backbone = CAE_Base_patch16_224_Async(pretrained,
                                              drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                              add_cls_token=cfg.MODEL.USE_CLS_TOKEN,
                                              num_async_interaction_stage=cfg.MODEL.BACKBONE.AI_LAYERS,
                                              depth=cfg.MODEL.BACKBONE.DEPTH
                                              )
        hidden_dim = backbone.embed_dim
        patch_start_index = 1

    else:
        raise NotImplementedError

    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)

    box_head = build_box_head(cfg, hidden_dim)

    stride = cfg.MODEL.BACKBONE.STRIDE
    feat_size_s = int(cfg.DATA.SEARCH.SIZE / stride)
    feat_size_t = int(cfg.DATA.TEMPLATE.SIZE / stride)
    if not training:
        cfg.MODEL.AUX_TEMPLATE = False
    model = LiteTrack(
        backbone,
        box_head,
        head_type=cfg.MODEL.HEAD.TYPE,
        search_feat_size=feat_size_s,
        template_feat_size=feat_size_t,
        add_target_token=cfg.MODEL.USE_CLS_TOKEN,
    )

    if 'LiteTrack' in cfg.MODEL.PRETRAIN_FILE and training:
        checkpoint = torch.load(cfg.MODEL.PRETRAIN_FILE, map_location="cpu")
        missing_keys, unexpected_keys = model.load_state_dict(
            checkpoint["net"], strict=False)
        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)

    return model

lib/models/litetrack/litetrack.py
TsingWei/LiteTrack
# Copyright (c) ByteDance, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
Mostly copy-paste from DINO and timm library:
https://github.com/facebookresearch/dino
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
"""


import math
import torch
import torch.nn as nn
import torch.nn.functional as F


from functools import partial
from timm.models.layers import to_2tuple


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    # if (mean < a - 2 * std) or (mean > b + 2 * std):
    #     warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
    #                   "The distribution of values may be incorrect.",
    #                   stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def drop_path(x, drop_prob: float = 0., training: bool = False):
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    # work with diff dim tensors, not just 2D ConvNets
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
    random_tensor = keep_prob + \
        torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., window_size=None):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        # add by wxd
        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        all_head_dim = head_dim * self.num_heads
        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.v_bias = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, len_t=None, x_rel_pos_bias=None, return_attention=True, mask=None):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(
                self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            attn = attn + x_rel_pos_bias.unsqueeze(0)
        if mask is not None:
            attn = attn.masked_fill(~mask, float('-inf'),)
        attn_nosf = attn
        attn = attn.softmax(dim=-1)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        if return_attention:
            if mask is not None:
                attn = attn[:, :, :-1, :-1]
            return x, attn_nosf
        else:
            return x


class MixedCrossAttention(Attention):
    def forward(self, x, len_t):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(
                self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q_s = q[:, :, len_t:]

        attn = (q_s @ k.transpose(-2, -1)) * self.scale

        attn = self.attn_drop(attn)

        attn_s2ts = attn
        attn_s2ts = attn_s2ts.softmax(dim=-1)

        x_t = x[:, :len_t]
        x_s = (attn_s2ts @ v).transpose(1, 2).reshape(B, N-len_t, C)
        x = torch.cat((x_t, x_s), dim=1)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., seperate_self_attn=False,
                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, window_size=None, init_values=0,
                 add_cls_token=False):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.add_cls_token = add_cls_token
        self.drop_path = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
                       act_layer=act_layer, drop=drop)
        self.attn = MixedCrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                                        attn_drop=attn_drop, proj_drop=drop, window_size=window_size)

        if init_values > 0:
            self.gamma_1 = nn.Parameter(
                init_values * torch.ones((dim)), requires_grad=True)
            self.gamma_2 = nn.Parameter(
                init_values * torch.ones((dim)), requires_grad=True)
        else:
            self.gamma_1, self.gamma_2 = None, None

    def forward_old(self, x, len_t):
        y = self.attn(self.norm1(x), len_t)
        if self.gamma_1 is None:
            x = x + self.drop_path(y)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.gamma_1 * y)
            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        return x

    def forward(self, x, len_t):
        # Here we keep the template part unchanged
        x_t = x[:, :len_t]
        x_s = x[:, len_t:]

        y= self.attn(self.norm1(
            torch.cat((x_t, x_s), dim=1)), len_t)
        y_s = y[:, len_t:]

        if self.gamma_1 is None:
            x = x + self.drop_path(y)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x_s = x_s + self.drop_path(self.gamma_1 * y_s)
            x_s = x_s + self.drop_path(self.gamma_2 *
                                       self.mlp(self.norm2(x_s)))
            x = torch.cat((x_t, x_s), dim=1)
        return x


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, img_size=[224, 224], patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.num_patches_w = img_size[0] // patch_size
        self.num_patches_h = img_size[1] // patch_size

        num_patches = self.num_patches_w * self.num_patches_h
        self.patch_shape = (img_size[0] // patch_size,
                            img_size[1] // patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        # print("##############patch here!!!")
        self.proj = nn.Conv2d(in_chans, embed_dim,
                              kernel_size=patch_size, stride=patch_size)
        self.flatten = True

    def forward(self, x, mask=None):
        # B, C, H, W = x.shape
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        # x = self.norm(x)
        return x


class VIT_Backbone(nn.Module):
    def __init__(self,
                 img_size=[224, 224],
                 patch_size=16,
                 embed_dim=768,
                 in_chans=3,
                 frozen_stages=-1,
                 #  out_indices=[3, 5, 7, 11],
                 out_with_norm=False,
                 use_checkpoint=False,
                 num_async_interaction_stage=3,
                 add_cls_token=False,
                 add_target_token=True,

                 num_classes=0,
                 depth=12,
                 num_heads=12,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 norm_layer=partial(nn.LayerNorm, eps=1e-6),
                 return_all_tokens=False,
                 init_values=0,
                 use_sincos_pos_emb=False,
                 use_abs_pos_emb=False,
                 use_rel_pos_bias=False,
                 seperate_loc=None,
                 act_layer=None,
                 ce_loc=None,
                 ce_keep_ratio=None,
                 use_mean_pooling=False,
                 masked_im_modeling=False,
                 ):

        super().__init__()

        self.num_features = self.embed_dim = embed_dim
        self.return_all_tokens = return_all_tokens
        self.use_abs_pos_emb = use_abs_pos_emb
        self.use_sincos_pos_emb = use_sincos_pos_emb
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        if use_abs_pos_emb:
            self.pos_embed = nn.Parameter(
                torch.zeros(1, num_patches + 1, embed_dim))
        else:
            self.pos_embed = None

        self.pos_drop = nn.Dropout(p=drop_rate)
        act_layer = act_layer or nn.GELU
        self.use_rel_pos_bias = use_rel_pos_bias

        if self.use_rel_pos_bias:
            print("=================use RelativePositionBias===================")
            window_size = self.patch_embed.patch_shape
            self.window_size = window_size
            self.num_relative_distance = (
                2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
            self.relative_position_bias_table = nn.Parameter(
                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
            # cls to token & token 2 cls & cls to cls

        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        # stochastic depth decay rule
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]

        blocks = []

        for i in range(depth):
            blocks.append(Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[
                    i], norm_layer=norm_layer,
                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None))
        self.blocks = nn.ModuleList(blocks)

        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
        if use_abs_pos_emb:
            trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        self.apply(self._init_weights)

        # masked image modeling
        self.masked_im_modeling = masked_im_modeling
        if masked_im_modeling:
            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))

        # support non-square image as input
        if len(img_size) == 1:
            img_size = img_size * 2
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches
        if self.use_abs_pos_emb:
            self.pos_embed = nn.Parameter(
                torch.zeros(1, num_patches + 1, embed_dim))
        elif self.use_sincos_pos_emb:
            self.pos_embed = self.build_2d_sincos_position_embedding(embed_dim)
        else:
            self.pos_embed = None

        self.add_target_token = add_target_token
        self.patch_size = patch_size
        self.frozen_stages = frozen_stages
        self.use_checkpoint = use_checkpoint
        self.num_FE_layers = depth - num_async_interaction_stage
        self.depth = depth
        # self.num_async_interaction_stage = num_async_interaction_stage

        self.pos_embed_z = None
        self.pos_embed_x = None

        if not out_with_norm:
            self.norm = nn.Identity()

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def finetune_track(self, cfg, patch_start_index=1):
        # patch_start_index: For skipping the cls embedding
        search_size = to_2tuple(cfg.DATA.SEARCH.SIZE)
        template_size = to_2tuple(cfg.DATA.TEMPLATE.SIZE)
        new_patch_size = cfg.MODEL.BACKBONE.STRIDE

        # resize patch embedding
        if new_patch_size != self.patch_size:
            print(
                'Inconsistent Patch Size With The Pretrained Weights, Interpolate The Weight!')
            old_patch_embed = {}
            for name, param in self.patch_embed.named_parameters():
                if 'weight' in name:
                    param = nn.functional.interpolate(param, size=(new_patch_size, new_patch_size),
                                                      mode='bicubic', align_corners=False)
                    param = nn.Parameter(param)
                old_patch_embed[name] = param
            # self.patch_embed = PatchEmbed(img_size=self.img_size, patch_size=new_patch_size, in_chans=3,
            #                               embed_dim=self.embed_dim)
            self.patch_embed.proj.bias = old_patch_embed['proj.bias']
            self.patch_embed.proj.weight = old_patch_embed['proj.weight']

        # Pos embedding
        patch_pos_embed = self.pos_embed[:, patch_start_index:, :]
        patch_pos_embed = patch_pos_embed.transpose(1, 2)  # ?
        B, E, Q = patch_pos_embed.shape
        P_H, P_W = self.patch_embed.patch_shape[0], self.patch_embed.patch_shape[1]
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)

        # Pos embedding for search region
        H, W = search_size
        new_P_H_S, new_P_W_S = H // new_patch_size, W // new_patch_size
        search_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H_S, new_P_W_S), mode='bicubic',
                                                           align_corners=False)
        search_patch_pos_embed = search_patch_pos_embed.flatten(
            2).transpose(1, 2)

        # Pos embedding for template region
        H, W = template_size
        new_P_H_T, new_P_W_T = H // new_patch_size, W // new_patch_size
        template_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H_T, new_P_W_T), mode='bicubic',
                                                             align_corners=False)
        template_patch_pos_embed = template_patch_pos_embed.flatten(
            2).transpose(1, 2)

        self.pos_embed_z = nn.Parameter(template_patch_pos_embed)
        self.pos_embed_x = nn.Parameter(search_patch_pos_embed)



    def train(self, mode=True):
        """Convert the model into training mode while keep layers freezed."""
        super(VIT_Backbone, self).train(mode)
        self._freeze_stages()

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.patch_embed.eval()
            for param in self.patch_embed.parameters():
                param.requires_grad = False
            self.cls_token.requires_grad = False
            if self.pos_embed is not None and self.use_sincos_pos_emb == True:
                self.pos_embed.requires_grad = False
            self.pos_drop.eval()

        for i in range(1, self.frozen_stages + 1):

            if i == len(self.blocks):
                norm_layer = getattr(self, 'norm')  # f'norm{i-1}')
                norm_layer.eval()
                for param in norm_layer.parameters():
                    param.requires_grad = False

            m = self.blocks[i - 1]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

    def forward_zx(self, z, x, target_token=None):
        x = self.patch_embed(x)
        x += self.pos_embed_x
        x = self.pos_drop(x)

        len_z = self.pos_embed_z.shape[1]

        if self.add_target_token and target_token is not None:
            z = torch.cat([target_token, z], dim=1)

        for i in range(self.num_FE_layers):
            x = self.blocks[i](x, 0)

        if self.add_target_token and target_token is not None:
            z = z[:,1:]

        x = torch.cat((z, x),dim=1)
        for i in range(self.num_FE_layers, self.depth):
            x = self.blocks[i](x, len_z)

        x = self.norm(x)
        return x

    def forward_z(self,
                  z,
                  target_token=None,
                  ):

        z = self.patch_embed(z)
        z += self.pos_embed_z
        z = self.pos_drop(z)

        if self.add_target_token and target_token is not None:
            z = torch.cat([target_token, z], dim=1)

        for i, blk in enumerate(self.blocks):
            z = blk(z, 0)

        if self.add_target_token and target_token is not None:
            z = z[:,1:]

        return z

    def forward(self, z, x=None,
                target_token=None,
                mode=None
                ):
        if mode is None:
            z_feat = self.forward_z(z)
            x_feat = self.forward_zx(z_feat, x, target_token=target_token)
        elif mode == 'x':
            x_feat = self.forward_zx(z, x)
        elif mode == 'z':
            x_feat = self.forward_z(z, target_token=target_token)


        return x_feat


def _create_ViT_CAE(pretrained=False, **kwargs) -> VIT_Backbone:
    model = VIT_Backbone(**kwargs)

    if pretrained:
        if 'npz' in pretrained:
            model.load_pretrained(pretrained, prefix='')
        else:
            state_dict = torch.load(
                'pretrained_models/cae_base.pth', map_location="cpu")
            # new_ckpt = OrderedDict()
            state_dict = state_dict['model']
            if sorted(list(state_dict.keys()))[0].startswith('encoder'):
                state_dict = {k.replace(
                    'encoder.', ''): v for k, v in state_dict.items() if k.startswith('encoder.')}
            missing_keys, unexpected_keys = model.load_state_dict(
                state_dict, strict=False)
            print("Unexpected keys:")
            print(unexpected_keys)
            print("Missing keys:")
            print(missing_keys)
    return model


def CAE_Base_patch16_224_Async(pretrained=False, **kwargs) -> VIT_Backbone:
    model_kwargs = dict(mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
                        patch_size=16, embed_dim=768, num_heads=12,
                        drop_rate=0., attn_drop_rate=0., init_values=0.1,
                        use_abs_pos_emb=True, **kwargs)
    model = _create_ViT_CAE(pretrained=pretrained, **model_kwargs)
    # model.adpat4track()
    return model

lib/models/litetrack/vit_cae_async.py
TsingWei/LiteTrack
from .ostrack import build_ostrack

lib/models/ostrack/__init__.py
TsingWei/LiteTrack
from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.vision_transformer import resize_pos_embed
from timm.models.layers import DropPath, to_2tuple, trunc_normal_

from lib.models.layers.patch_embed import PatchEmbed
from lib.models.ostrack.utils import combine_tokens, recover_tokens


class BaseBackbone(nn.Module):
    def __init__(self):
        super().__init__()

        # for original ViT
        self.pos_embed = None
        self.img_size = [224, 224]
        self.patch_size = 16
        self.embed_dim = 384

        self.cat_mode = 'direct'

        self.pos_embed_z = None
        self.pos_embed_x = None

        self.template_segment_pos_embed = None
        self.search_segment_pos_embed = None

        self.return_inter = False
        self.return_stage = [2, 5, 8, 11]

        self.add_cls_token = False
        self.add_sep_seg = False

    def finetune_track(self, cfg, patch_start_index=1):

        search_size = to_2tuple(cfg.DATA.SEARCH.SIZE)
        template_size = to_2tuple(cfg.DATA.TEMPLATE.SIZE)
        new_patch_size = cfg.MODEL.BACKBONE.STRIDE

        self.cat_mode = cfg.MODEL.BACKBONE.CAT_MODE
        self.return_inter = cfg.MODEL.RETURN_INTER
        self.add_sep_seg = cfg.MODEL.BACKBONE.SEP_SEG

        # resize patch embedding
        if new_patch_size != self.patch_size:
            print('Inconsistent Patch Size With The Pretrained Weights, Interpolate The Weight!')
            old_patch_embed = {}
            for name, param in self.patch_embed.named_parameters():
                if 'weight' in name:
                    param = nn.functional.interpolate(param, size=(new_patch_size, new_patch_size),
                                                      mode='bicubic', align_corners=False)
                    param = nn.Parameter(param)
                old_patch_embed[name] = param
            self.patch_embed = PatchEmbed(img_size=self.img_size, patch_size=new_patch_size, in_chans=3,
                                          embed_dim=self.embed_dim)
            self.patch_embed.proj.bias = old_patch_embed['proj.bias']
            self.patch_embed.proj.weight = old_patch_embed['proj.weight']

        # for patch embedding
        patch_pos_embed = self.pos_embed[:, patch_start_index:, :]
        patch_pos_embed = patch_pos_embed.transpose(1, 2)
        B, E, Q = patch_pos_embed.shape
        P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)

        # for search region
        H, W = search_size
        new_P_H, new_P_W = H // new_patch_size, W // new_patch_size
        search_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H, new_P_W), mode='bicubic',
                                                           align_corners=False)
        search_patch_pos_embed = search_patch_pos_embed.flatten(2).transpose(1, 2)

        # for template region
        H, W = template_size
        new_P_H, new_P_W = H // new_patch_size, W // new_patch_size
        template_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H, new_P_W), mode='bicubic',
                                                             align_corners=False)
        template_patch_pos_embed = template_patch_pos_embed.flatten(2).transpose(1, 2)

        self.pos_embed_z = nn.Parameter(template_patch_pos_embed)
        self.pos_embed_x = nn.Parameter(search_patch_pos_embed)

        # for cls token (keep it but not used)
        if self.add_cls_token and patch_start_index > 0:
            cls_pos_embed = self.pos_embed[:, 0:1, :]
            self.cls_pos_embed = nn.Parameter(cls_pos_embed)

        # separate token and segment token
        if self.add_sep_seg:
            self.template_segment_pos_embed = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
            self.template_segment_pos_embed = trunc_normal_(self.template_segment_pos_embed, std=.02)
            self.search_segment_pos_embed = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
            self.search_segment_pos_embed = trunc_normal_(self.search_segment_pos_embed, std=.02)

        # self.cls_token = None
        # self.pos_embed = None

        if self.return_inter:
            for i_layer in self.fpn_stage:
                if i_layer != 11:
                    norm_layer = partial(nn.LayerNorm, eps=1e-6)
                    layer = norm_layer(self.embed_dim)
                    layer_name = f'norm{i_layer}'
                    self.add_module(layer_name, layer)

    def forward_features(self, z, x):
        B, H, W = x.shape[0], x.shape[2], x.shape[3]

        x = self.patch_embed(x)
        z = self.patch_embed(z)

        if self.add_cls_token:
            cls_tokens = self.cls_token.expand(B, -1, -1)
            cls_tokens = cls_tokens + self.cls_pos_embed

        z += self.pos_embed_z
        x += self.pos_embed_x

        if self.add_sep_seg:
            x += self.search_segment_pos_embed
            z += self.template_segment_pos_embed

        x = combine_tokens(z, x, mode=self.cat_mode)
        if self.add_cls_token:
            x = torch.cat([cls_tokens, x], dim=1)

        x = self.pos_drop(x)

        for i, blk in enumerate(self.blocks):
            x = blk(x)

        lens_z = self.pos_embed_z.shape[1]
        lens_x = self.pos_embed_x.shape[1]
        x = recover_tokens(x, lens_z, lens_x, mode=self.cat_mode)

        aux_dict = {"attn": None}
        return self.norm(x), aux_dict

    def forward(self, z, x, **kwargs):
        """
        Joint feature extraction and relation modeling for the basic ViT backbone.
        Args:
            z (torch.Tensor): template feature, [B, C, H_z, W_z]
            x (torch.Tensor): search region feature, [B, C, H_x, W_x]

        Returns:
            x (torch.Tensor): merged template and search region feature, [B, L_z+L_x, C]
            attn : None
        """
        x, aux_dict = self.forward_features(z, x,)

        return x, aux_dict

lib/models/ostrack/base_backbone.py
TsingWei/LiteTrack
"""
Basic OSTrack model.
"""
import math
import os
from typing import List

import torch
from torch import nn
from torch.nn.modules.transformer import _get_clones

from lib.models.layers.head import build_box_head
from lib.models.ostrack.vit import vit_base_patch16_224
from lib.models.ostrack.vit_ce import vit_large_patch16_224_ce, vit_base_patch16_224_ce
from lib.models.ostrack.vit_cae import CAE_Base_patch16_224
from lib.utils.box_ops import box_xyxy_to_cxcywh



class ContrastiveEmbed(nn.Module):
    '''
    A classification head using dynamic objects as refference.
    '''
    def __init__(self, max_obj_len=256, num_classes=1, in_dim=768, out_dim=256):
        """
        Args:
            max_obj_len: max length of obj.
        """
        super().__init__()
        self.max_obj_len = max_obj_len
        self.num_classes = num_classes
        self.transform = nn.Linear(in_dim, out_dim, bias=False)

    def forward(self, x, obj):
        """_summary_

        Args:
            x (tensor): unknown obj
            obj (tensor): known obj
        Returns:
            _type_: _description_
        """
        y = obj
        # text_token_mask = text_dict["text_token_mask"]

        res = self.transform(x) @ self.transform(y).transpose(-1, -2)
        # res.masked_fill_(~text_token_mask[:, None, :], float("-inf"))

        # padding to max_text_len
        # new_res = torch.full((*res.shape[:-1], self.max_obj_len), float("-inf"), device=res.device, dtype=torch.float16)
        # new_res[..., : res.shape[-1]] = res

        return res


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class OSTrack(nn.Module):
    """ This is the base class for OSTrack """

    def __init__(self, transformer, box_head, aux_loss=False, head_type="CORNER", aux_template_loss=False,
                    search_feat_size=16, template_feat_size=8, add_target_token=False):
        """ Initializes the model.
        Parameters:
            transformer: torch module of the transformer architecture.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.backbone = transformer
        self.box_head = box_head

        self.aux_loss = aux_loss
        self.head_type = head_type
        # feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
        # if head_type == "CORNER" or head_type == "CENTER":
        self.feat_size_s = int(search_feat_size)
        self.feat_len_s = int(search_feat_size ** 2)
        self.feat_size_t = int(template_feat_size)
        self.feat_len_t = int(template_feat_size ** 2)

        if self.aux_loss:
            self.box_head = _get_clones(self.box_head, 6)
        self.aux_template_loss = aux_template_loss
        self.add_target_token = add_target_token
        if self.add_target_token:
            self.target_token_embed = Mlp(4, out_features=self.backbone.embed_dim)
        self.onnx = True
        # self.neck = MLP(768,768,256,3)
        # self.prpool = PrRoIPool(1)
        # self.label_enc = nn.Embedding(4 + 1, 768)
        # self.constrative_head = ContrastiveEmbed()

    def forward_z():
        pass

    def forward_x():
        pass
            
    def forward(self, template: torch.Tensor,
                search: torch.Tensor,
                template_bb=None,
                ce_template_mask=None,
                ce_keep_rate=None,
                return_last_attn=False,
                search_bb=None,
                gt_cls_map = None,
                threshold=0,
                gt_cls_map_T=None,
                mask_z=None,
                mask_s=None,
                mode=None,
                temporal_feat=None,
                head_only=False,
                ):
        if False:
            feat, _ = self.backbone(z=template, x=search,
                                ce_template_mask=ce_template_mask,
                                ce_keep_rate=ce_keep_rate,
                                template_bb=template_bb,
                                mode=mode)
            if mode=='z':
                return feat
            return self.forward_head(feat, gt_cls_map,
                                 gt_score_map_T=gt_cls_map_T, 
                                 th=threshold, 
                                 mask_z=mask_z,
                                 mask_s=mask_s,
                                 template_bb=template_bb,
                                 search_bb=search_bb)
            

                
        # if self.add_target_token:
        #     target_token = self.target_token_embed(template_bb).unsqueeze(-2)
        #     # x, aux_dict = self.backbone(z=template, x=search,
        #     #                         ce_template_mask=ce_template_mask,
        #     #                         ce_keep_rate=ce_keep_rate,
        #     #                         return_last_attn=return_last_attn,
        #     #                         template_bb=template_bb,
        #     #                         target_token=target_token)
        # else:

        if self.add_target_token:
            target_token = self.target_token_embed(template_bb).unsqueeze(-2)
            x, aux_dict = self.backbone(z=template, x=search,
                            ce_template_mask=ce_template_mask,
                            ce_keep_rate=ce_keep_rate,
                            target_token=target_token
                            # return_last_attn=return_last_attn,
                            )
        else:
            x, aux_dict = self.backbone(z=template, x=search,
                            ce_template_mask=ce_template_mask,
                            ce_keep_rate=ce_keep_rate,
                            # target_token=target_token
                            # return_last_attn=return_last_attn,
                            )
            

        # Forward head
        feat_last = x
        if isinstance(x, list):
            feat_last = x[-1]
        attn = aux_dict['attn']
        out = {}
        # return
        # size_token = x[:, -1:]
        # x = x[:, :-1]
        # guess_size = self.size_guessor(size_token).squeeze_(1).sigmoid()
        # feat_search = x[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
        # pooled_feat = self.prpool(feat_search.permute((0, 2, 1)).contiguous().reshape(32,-1,16,16), template_bb)
        # out = {}

        


        out = self.forward_head(x, gt_cls_map,
                                gt_score_map_T=gt_cls_map_T, 
                                th=threshold, wh=None, attn=attn, 
                                mask_z=mask_z,
                                mask_s=mask_s,
                                template_bb=template_bb,
                                search_bb=search_bb,
                                )
        if self.onnx:
            return out['score_map'], out['size_map'], out['offset_map']

        out.update(aux_dict)
        # out['feat_search'] = feat_search
        # out['guess_size'] = guess_size
        return out

    def forward_head(self, cat_feature,
                    gt_score_map=None,
                    th=0, 
                    gt_score_map_T=None,
                    wh=None, attn=None, 
                    mask_z=None, 
                    mask_s=None, 
                    template_bb=None, 
                    search_bb=None,
                    **kwargs,
                    ):
        """
        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
        """
        if self.aux_template_loss:
            # TODO if cls token?
            enc_opt_T = cat_feature[:, :self.feat_len_t]  # encoder output for the search region (B, HW, C)
            opt_T = (enc_opt_T.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt_T.size()
            opt_feat_T = opt_T.view(-1, C, self.feat_size_t, self.feat_size_t)
        # enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
        # opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
        # bs, Nq, C, HW = opt.size()
        # opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)

        if self.head_type == "CORNER":
            # run the corner head
            pred_box, score_map = self.box_head(opt_feat, True)
            outputs_coord = box_xyxy_to_cxcywh(pred_box)
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map,
                   }
            return out
        elif self.head_type == "CENTER":
            enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)
            # run the center head
            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat, self.feat_size_s, gt_score_map)
            # outputs_coord = box_xyxy_to_cxcywh(bbox)
            outputs_coord = bbox
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map_ctr,
                   'size_map': size_map,
                   'offset_map': offset_map}
            if self.aux_template_loss:
                # shared center head
                score_map_ctr_T, bbox_T, size_map_T, offset_map_T = self.box_head(opt_feat_T, self.feat_size_t,)
                outputs_coord_T = bbox_T
                outputs_coord_new_T = outputs_coord_T.view(bs, Nq, 4)
                out['pred_boxes_T'] = outputs_coord_new_T
                out['score_map_T'] = score_map_ctr_T
                out['size_map_T'] = size_map_T
                out['offset_map_T'] = offset_map_T
        
            return out
        elif self.head_type == "GFL":
            # run the head

            enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)
            box_best = None
            score_map_ctr, bboxes, ltrb_distribution, box_best, qs = self.box_head(opt_feat, self.feat_size_s, gt_score_map, th=th)
            outputs_coord = box_xyxy_to_cxcywh(bboxes)
            # bboxes = box_xyxy_to_cxcywh(bboxes)

            # Pooling
            # batch_index = torch.arange(bs, dtype=torch.float32).view(-1, 1).to(template_bb.device)
            # box_z = template_bb * 16
            # pooled_z = self.prpool(feat_t, torch.cat((batch_index, box_z), dim=1)).flatten(2).permute((0, 2, 1))
            # box_x = box_best * 16
            # pooled_x = self.prpool(feat_s, torch.cat((batch_index, box_x), dim=1)).flatten(2).permute((0, 2, 1))
            # box_x_GT = search_bb *16
            # pooled_x_GT = self.prpool(feat_s, torch.cat((batch_index, box_x_GT), dim=1)).flatten(2).permute((0, 2, 1))
            # corr = self.constrative_head(pooled_z, pooled_x)
            # corr_GT = self.constrative_head(pooled_z, pooled_x_GT)

            # outputs_coord = bbox
            # outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord,
                #    'score_map_c': score_map_c,
                   'score_map': score_map_ctr,
                   'ltrb_dis': ltrb_distribution,
                   'box_best': box_best,
                   'boxes': bboxes,
                   'qs': qs,
                #    'corr': corr,
                #    'corr2GT': corr_GT,
                   }
            if self.aux_template_loss:
                # shared center head
                box_best_T = None
                score_map_ctr_T, bboxes_T, ltrb_distribution_T, box_best_T = self.box_head(opt_feat_T, self.feat_size_t, gt_score_map_T, th=th)
                outputs_coord_T = box_xyxy_to_cxcywh(bboxes_T)
                # outputs_coord_new_T = outputs_coord_T.view(bs, Nq, 4)
                out['pred_boxes_T'] = outputs_coord_T
                out['score_map_T'] = score_map_ctr_T
                out['ltrb_dis_T'] = ltrb_distribution_T
                out['box_best_T'] = box_best_T
                out['boxes_T'] = bboxes_T
            return out
        elif self.head_type == "PyGFL":
            # run the head
            box_best = None
            score_map_ctr, bboxes, ltrb_distribution, box_best = self.box_head(cat_feature, self.feat_size_s, gt_score_map, th=th)
            outputs_coord = box_xyxy_to_cxcywh(bboxes)
            # bboxes = box_xyxy_to_cxcywh(bboxes)

            # outputs_coord = bbox
            # outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord,
                   'score_map': score_map_ctr,
                   'ltrb_dis': ltrb_distribution,
                   'box_best': box_best,
                   'boxes': bboxes
                   }
            return out
        elif self.head_type == "REPCENTER":
            # run the head
            enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)
            box_best = None
            cls_map, bboxes, _, box_best = self.box_head(opt_feat, self.feat_size_s, gt_score_map, th=th)
            outputs_coord = bboxes
            # bboxes = box_xyxy_to_cxcywh(bboxes)

            # outputs_coord = bbox
            # outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord,
                   'score_map': cls_map,
                #    'ltrb_dis': ltrb_distribution,
                   'box_best': box_best,
                   'boxes': bboxes
                   }
            return out
        elif self.head_type == "REPPOINTS":
            # run the head
            enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            opt_feat = opt.view(-1, C, self.feat_size_s, self.feat_size_s)
            box_best = None
            # box_in = torch.cat((-wh*0.5, wh*0.5), -1) * self.feat_size_s 
            # box_in = box_in[...,None,None].repeat(1,1,16,16).cuda()
            cls_score, init_box, refine_box, init_pts, refine_pts, anchor_box = self.box_head(opt_feat, self.feat_size_s, gt_score_map, th=th, box_in=None)
            # xyxy
            out = { 'score_map': cls_score,
                    'init_box': init_box,
                    'refine_box': refine_box,
                    'init_pts': init_pts,
                    'refine_pts': refine_pts,
                    'anchor_box': anchor_box,
                   }
            return out
        elif self.head_type == "DECODER":

            # cdn
            

            # run the head
            # cat_feature = self.neck(cat_feature)

            enc_opt = cat_feature[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
            enc_opt_T = cat_feature[:, :self.feat_len_t].to(torch.float32)
            
            # pooled_feat = self.prpool(enc_opt.reshape(32,16,16,-1), )
            # 
            
            opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt.size()
            feat_s = opt.view(-1, C, self.feat_size_s, self.feat_size_s)
            feat_t = enc_opt_T.permute((0, 2, 1)).contiguous().reshape(bs,-1,self.feat_size_t,self.feat_size_t)
            

            box_best = None
            out_dict={}
            out_dict = self.box_head(feat_s, feat_t, attn, mask_z, mask_s, template_bb, search_bb, **kwargs,)
            # out_dict['template_tokens'] = template_feat
            # # xyxy
            # out = { 'score_map': cls_score,
            #         'init_box': init_box,
            #         'refine_box': refine_box,
            #         'init_pts': init_pts,
            #         'refine_pts': refine_pts,
            #         'anchor_box': anchor_box,
            #        }
            return out_dict
        
        else:
            raise NotImplementedError

class SimpleFPN(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()

        # # Up sampling 4X
        # self.fpn1 = nn.Sequential(
        #     nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
        #     Norm2d(embed_dim),
        #     nn.GELU(),
        #     nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
        # )

        # Up sampling 2X
        self.fpn2 = nn.Sequential(
            nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
        )

        # Keep 1X
        self.fpn3 = nn.Identity()

        # Down sampling 0.5X
        self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)

    
    def forward(self, x):
        B, C, H, W = x.shape
        xp = x.permute(0, 2, 1).reshape(B, -1, H, W)
        features = []
        ops = [ self.fpn2, self.fpn3, self.fpn4]
        for i in range(len(ops)):
            features.append(ops[i](xp))
        return features

def build_ostrack(cfg, training=True):
    current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
    pretrained_path = os.path.join(current_dir, '../../../pretrained_models')
    if cfg.MODEL.PRETRAIN_FILE and ('OSTrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
    else:
        pretrained = ''

    if cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224_CAE':
        backbone = CAE_Base_patch16_224(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                        ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                        ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                        CE=cfg.MODEL.BACKBONE.CE,
                                        add_cls_token = cfg.MODEL.USE_CLS_TOKEN,
                                        seperate_loc = cfg.MODEL.BACKBONE.SEP_LOC,
                                        )
        hidden_dim = backbone.embed_dim
        patch_start_index = 1
    
    elif cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224_ce':
        backbone = vit_base_patch16_224_ce(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                           ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                           ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                           )
        hidden_dim = backbone.embed_dim
        patch_start_index = 1
    elif cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224':
        backbone = vit_base_patch16_224(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                           )
        hidden_dim = backbone.embed_dim
        patch_start_index = 1


    elif cfg.MODEL.BACKBONE.TYPE == 'cae_fpn':
        backbone = cae_fpn(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                           ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                           ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                           )
        hidden_dim = backbone.embed_dim
        patch_start_index = 1

    elif cfg.MODEL.BACKBONE.TYPE == 'swin':
        backbone = SwinBackbone4Track(pretrained)
        hidden_dim = backbone.embed_dim
        patch_start_index = 0

    elif cfg.MODEL.BACKBONE.TYPE == 'vit_large_patch16_224_ce':
        backbone = vit_large_patch16_224_ce(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                            ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                            ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                            )

        hidden_dim = backbone.embed_dim
        patch_start_index = 1

    else:
        raise NotImplementedError

    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)

    box_head = build_box_head(cfg, hidden_dim)

    stride = cfg.MODEL.BACKBONE.STRIDE
    feat_size_s = int(cfg.DATA.SEARCH.SIZE / stride)
    feat_size_t = int(cfg.DATA.TEMPLATE.SIZE / stride)
    if not training:
        cfg.MODEL.AUX_TEMPLATE = False
    model = OSTrack(
        backbone,
        box_head,
        aux_loss=False,
        head_type=cfg.MODEL.HEAD.TYPE,
        search_feat_size=feat_size_s,
        template_feat_size=feat_size_t,
        aux_template_loss=cfg.MODEL.AUX_TEMPLATE,
        add_target_token=cfg.MODEL.USE_CLS_TOKEN,
    )

    if 'OSTrack' in cfg.MODEL.PRETRAIN_FILE and training:
        checkpoint = torch.load(cfg.MODEL.PRETRAIN_FILE, map_location="cpu")
        missing_keys, unexpected_keys = model.load_state_dict(checkpoint["net"], strict=False)
        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)

    return model

lib/models/ostrack/ostrack.py
TsingWei/LiteTrack
import math

import torch
import torch.nn.functional as F


def combine_tokens(template_tokens, search_tokens, mode='direct', return_res=False):
    # [B, HW, C]
    len_t = template_tokens.shape[1]
    len_s = search_tokens.shape[1]

    if mode == 'direct':
        merged_feature = torch.cat((template_tokens, search_tokens), dim=1)
    elif mode == 'template_central':
        central_pivot = len_s // 2
        first_half = search_tokens[:, :central_pivot, :]
        second_half = search_tokens[:, central_pivot:, :]
        merged_feature = torch.cat((first_half, template_tokens, second_half), dim=1)
    elif mode == 'partition':
        feat_size_s = int(math.sqrt(len_s))
        feat_size_t = int(math.sqrt(len_t))
        window_size = math.ceil(feat_size_t / 2.)
        # pad feature maps to multiples of window size
        B, _, C = template_tokens.shape
        H = W = feat_size_t
        template_tokens = template_tokens.view(B, H, W, C)
        pad_l = pad_b = pad_r = 0
        # pad_r = (window_size - W % window_size) % window_size
        pad_t = (window_size - H % window_size) % window_size
        template_tokens = F.pad(template_tokens, (0, 0, pad_l, pad_r, pad_t, pad_b))
        _, Hp, Wp, _ = template_tokens.shape
        template_tokens = template_tokens.view(B, Hp // window_size, window_size, W, C)
        template_tokens = torch.cat([template_tokens[:, 0, ...], template_tokens[:, 1, ...]], dim=2)
        _, Hc, Wc, _ = template_tokens.shape
        template_tokens = template_tokens.view(B, -1, C)
        merged_feature = torch.cat([template_tokens, search_tokens], dim=1)

        # calculate new h and w, which may be useful for SwinT or others
        merged_h, merged_w = feat_size_s + Hc, feat_size_s
        if return_res:
            return merged_feature, merged_h, merged_w

    else:
        raise NotImplementedError

    return merged_feature


def recover_tokens(merged_tokens, len_template_token, len_search_token, mode='direct'):
    if mode == 'direct':
        recovered_tokens = merged_tokens
    elif mode == 'template_central':
        central_pivot = len_search_token // 2
        len_remain = len_search_token - central_pivot
        len_half_and_t = central_pivot + len_template_token

        first_half = merged_tokens[:, :central_pivot, :]
        second_half = merged_tokens[:, -len_remain:, :]
        template_tokens = merged_tokens[:, central_pivot:len_half_and_t, :]

        recovered_tokens = torch.cat((template_tokens, first_half, second_half), dim=1)
    elif mode == 'partition':
        recovered_tokens = merged_tokens
    else:
        raise NotImplementedError

    return recovered_tokens


def window_partition(x, window_size: int):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size: int, H: int, W: int):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

class Norm2d(torch.nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.ln = torch.nn.LayerNorm(embed_dim, eps=1e-6)
    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        x = self.ln(x)
        x = x.permute(0, 3, 1, 2).contiguous()
        return x
lib/models/ostrack/utils.py
TsingWei/LiteTrack
""" Vision Transformer (ViT) in PyTorch
A PyTorch implement of Vision Transformers as described in:
'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale'
    - https://arxiv.org/abs/2010.11929
`How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers`
    - https://arxiv.org/abs/2106.10270
The official jax code is released and available at https://github.com/google-research/vision_transformer
DeiT model defs and weights from https://github.com/facebookresearch/deit,
paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877
Acknowledgments:
* The paper authors for releasing code and weights, thanks!
* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
for some einops/einsum fun
* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
Hacked together by / Copyright 2021 Ross Wightman

Modified by Botao Ye
"""
import math
import logging
from functools import partial
from collections import OrderedDict
from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F

from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
from timm.models.helpers import build_model_with_cfg, named_apply, adapt_input_conv
from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
from timm.models.registry import register_model

from lib.models.layers.patch_embed import PatchEmbed
from lib.models.ostrack.base_backbone import BaseBackbone


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, return_attention=False):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        if return_attention:
            return x, attn
        return x


class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, return_attention=False):
        if return_attention:
            feat, attn = self.attn(self.norm1(x), True)
            x = x + self.drop_path(feat)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
            return x, attn
        else:
            x = x + self.drop_path(self.attn(self.norm1(x)))
            x = x + self.drop_path(self.mlp(self.norm2(x)))
            return x


class VisionTransformer(BaseBackbone):
    """ Vision Transformer
    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
        - https://arxiv.org/abs/2010.11929
    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
        - https://arxiv.org/abs/2012.12877
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
                 act_layer=None, weight_init=''):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
            weight_init: (str): weight init scheme
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_tokens = 2 if distilled else 1
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        act_layer = act_layer or nn.GELU

        self.patch_embed = embed_layer(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.Sequential(*[
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
                attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        # # Representation layer
        # if representation_size and not distilled:
        #     self.num_features = representation_size
        #     self.pre_logits = nn.Sequential(OrderedDict([
        #         ('fc', nn.Linear(embed_dim, representation_size)),
        #         ('act', nn.Tanh())
        #     ]))
        # else:
        #     self.pre_logits = nn.Identity()
        #
        # # Classifier head(s)
        # self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
        # self.head_dist = None
        # if distilled:
        #     self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()

        self.init_weights(weight_init)

    def init_weights(self, mode=''):
        assert mode in ('jax', 'jax_nlhb', 'nlhb', '')
        head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0.
        trunc_normal_(self.pos_embed, std=.02)
        if self.dist_token is not None:
            trunc_normal_(self.dist_token, std=.02)
        if mode.startswith('jax'):
            # leave cls token as zeros to match jax impl
            named_apply(partial(_init_vit_weights, head_bias=head_bias, jax_impl=True), self)
        else:
            trunc_normal_(self.cls_token, std=.02)
            self.apply(_init_vit_weights)

    def _init_weights(self, m):
        # this fn left here for compat with downstream users
        _init_vit_weights(m)

    @torch.jit.ignore()
    def load_pretrained(self, checkpoint_path, prefix=''):
        _load_weights(self, checkpoint_path, prefix)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token', 'dist_token'}

    def get_classifier(self):
        if self.dist_token is None:
            return self.head
        else:
            return self.head, self.head_dist

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
        if self.num_tokens == 2:
            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()


def _init_vit_weights(module: nn.Module, name: str = '', head_bias: float = 0., jax_impl: bool = False):
    """ ViT weight initialization
    * When called without n, head_bias, jax_impl args it will behave exactly the same
      as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).
    * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl
    """
    if isinstance(module, nn.Linear):
        if name.startswith('head'):
            nn.init.zeros_(module.weight)
            nn.init.constant_(module.bias, head_bias)
        elif name.startswith('pre_logits'):
            lecun_normal_(module.weight)
            nn.init.zeros_(module.bias)
        else:
            if jax_impl:
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    if 'mlp' in name:
                        nn.init.normal_(module.bias, std=1e-6)
                    else:
                        nn.init.zeros_(module.bias)
            else:
                trunc_normal_(module.weight, std=.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
    elif jax_impl and isinstance(module, nn.Conv2d):
        # NOTE conv was left to pytorch default in my original init
        lecun_normal_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
        nn.init.zeros_(module.bias)
        nn.init.ones_(module.weight)


@torch.no_grad()
def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str = ''):
    """ Load weights from .npz checkpoints for official Google Brain Flax implementation
    """
    import numpy as np

    def _n2p(w, t=True):
        if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:
            w = w.flatten()
        if t:
            if w.ndim == 4:
                w = w.transpose([3, 2, 0, 1])
            elif w.ndim == 3:
                w = w.transpose([2, 0, 1])
            elif w.ndim == 2:
                w = w.transpose([1, 0])
        return torch.from_numpy(w)

    w = np.load(checkpoint_path)
    if not prefix and 'opt/target/embedding/kernel' in w:
        prefix = 'opt/target/'

    if hasattr(model.patch_embed, 'backbone'):
        # hybrid
        backbone = model.patch_embed.backbone
        stem_only = not hasattr(backbone, 'stem')
        stem = backbone if stem_only else backbone.stem
        stem.conv.weight.copy_(adapt_input_conv(stem.conv.weight.shape[1], _n2p(w[f'{prefix}conv_root/kernel'])))
        stem.norm.weight.copy_(_n2p(w[f'{prefix}gn_root/scale']))
        stem.norm.bias.copy_(_n2p(w[f'{prefix}gn_root/bias']))
        if not stem_only:
            for i, stage in enumerate(backbone.stages):
                for j, block in enumerate(stage.blocks):
                    bp = f'{prefix}block{i + 1}/unit{j + 1}/'
                    for r in range(3):
                        getattr(block, f'conv{r + 1}').weight.copy_(_n2p(w[f'{bp}conv{r + 1}/kernel']))
                        getattr(block, f'norm{r + 1}').weight.copy_(_n2p(w[f'{bp}gn{r + 1}/scale']))
                        getattr(block, f'norm{r + 1}').bias.copy_(_n2p(w[f'{bp}gn{r + 1}/bias']))
                    if block.downsample is not None:
                        block.downsample.conv.weight.copy_(_n2p(w[f'{bp}conv_proj/kernel']))
                        block.downsample.norm.weight.copy_(_n2p(w[f'{bp}gn_proj/scale']))
                        block.downsample.norm.bias.copy_(_n2p(w[f'{bp}gn_proj/bias']))
        embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])
    else:
        embed_conv_w = adapt_input_conv(
            model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))
    model.patch_embed.proj.weight.copy_(embed_conv_w)
    model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
    model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
    pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)
    if pos_embed_w.shape != model.pos_embed.shape:
        pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights
            pos_embed_w, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)
    model.pos_embed.copy_(pos_embed_w)
    model.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))
    model.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))
    if isinstance(model.head, nn.Linear) and model.head.bias.shape[0] == w[f'{prefix}head/bias'].shape[-1]:
        model.head.weight.copy_(_n2p(w[f'{prefix}head/kernel']))
        model.head.bias.copy_(_n2p(w[f'{prefix}head/bias']))
    if isinstance(getattr(model.pre_logits, 'fc', None), nn.Linear) and f'{prefix}pre_logits/bias' in w:
        model.pre_logits.fc.weight.copy_(_n2p(w[f'{prefix}pre_logits/kernel']))
        model.pre_logits.fc.bias.copy_(_n2p(w[f'{prefix}pre_logits/bias']))
    for i, block in enumerate(model.blocks.children()):
        block_prefix = f'{prefix}Transformer/encoderblock_{i}/'
        mha_prefix = block_prefix + 'MultiHeadDotProductAttention_1/'
        block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
        block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
        block.attn.qkv.weight.copy_(torch.cat([
            _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))
        block.attn.qkv.bias.copy_(torch.cat([
            _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))
        block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
        block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
        for r in range(2):
            getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/kernel']))
            getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/bias']))
        block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/scale']))
        block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/bias']))


def resize_pos_embed(posemb, posemb_new, num_tokens=1, gs_new=()):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    print('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    if num_tokens:
        posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:]
        ntok_new -= num_tokens
    else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
    gs_old = int(math.sqrt(len(posemb_grid)))
    if not len(gs_new):  # backwards compatibility
        gs_new = [int(math.sqrt(ntok_new))] * 2
    assert len(gs_new) >= 2
    print('Position embedding grid-size from %s to %s', [gs_old, gs_old], gs_new)
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = F.interpolate(posemb_grid, size=gs_new, mode='bilinear')
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new[0] * gs_new[1], -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return posemb


def checkpoint_filter_fn(state_dict, model):
    """ convert patch embedding weight from manual patchify + linear proj to conv"""
    out_dict = {}
    if 'model' in state_dict:
        # For deit models
        state_dict = state_dict['model']
    for k, v in state_dict.items():
        if 'patch_embed.proj.weight' in k and len(v.shape) < 4:
            # For old models that I trained prior to conv based patchification
            O, I, H, W = model.patch_embed.proj.weight.shape
            v = v.reshape(O, -1, H, W)
        elif k == 'pos_embed' and v.shape != model.pos_embed.shape:
            # To resize pos embedding when using model at different size from pretrained weights
            v = resize_pos_embed(
                v, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)
        out_dict[k] = v
    return out_dict


def _create_vision_transformer(variant, pretrained=False, default_cfg=None, **kwargs):
    if kwargs.get('features_only', None):
        raise RuntimeError('features_only not implemented for Vision Transformer models.')

    model = VisionTransformer(**kwargs)

    if pretrained:
        if 'npz' in pretrained:
            model.load_pretrained(pretrained, prefix='')
        else:
            checkpoint = torch.load(pretrained, map_location="cpu")
            missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False)
            print('Load pretrained model from: ' + pretrained)

    return model


def vit_base_patch16_224(pretrained=False, **kwargs):
    """
    ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    """
    model_kwargs = dict(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
    model = _create_vision_transformer('vit_base_patch16_224_in21k', pretrained=pretrained, **model_kwargs)
    return model

lib/models/ostrack/vit.py
TsingWei/LiteTrack
# Copyright (c) ByteDance, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
Mostly copy-paste from DINO and timm library:
https://github.com/facebookresearch/dino
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
"""

import imp
import os
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from .utils import combine_tokens, recover_tokens
from functools import partial
from timm.models.layers import to_2tuple


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    # if (mean < a - 2 * std) or (mean > b + 2 * std):
    #     warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
    #                   "The distribution of values may be incorrect.",
    #                   stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor

def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

def drop_path(x, drop_prob: float = 0., training: bool = False):
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., window_size=None):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        #self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        ####add by wxd
        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        all_head_dim = head_dim * self.num_heads
        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.v_bias = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, len_t=None, x_rel_pos_bias = None, return_attention = True, mask=None):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            attn = attn + x_rel_pos_bias.unsqueeze(0)
        if mask is not None:
            attn = attn.masked_fill(~mask, float('-inf'),)
        attn_nosf = attn
        attn = attn.softmax(dim=-1)
        
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        if return_attention:
            if mask is not None:
                attn = attn[:,:,:-1,:-1]
            return x, attn_nosf
        else:
            return x

class CenterAttention(Attention):
    # def __init__(self, **kwargs):
    #     super().__init__(**kwargs)
    #     self.center_attn_mask = 
    def forward(self, x, len_t=None, x_rel_pos_bias=None, return_attention=True, mask=None):
        
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            attn = attn + x_rel_pos_bias.unsqueeze(0)
        if mask is not None:
            attn = attn.masked_fill(~mask, float('-inf'),)
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        if return_attention:
            return x, attn
        else:
            return x

class SingleCrossAttention(Attention):
    def forward(self, x, len_t=None, x_rel_pos_bias = None, return_attention = True):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q_t, k_t, v_t = q[:, :, :len_t], k[:, :, :len_t], v[:, :, :len_t]
        q_s, k_s, v_s = q[:, :, len_t:], k[:, :, len_t:], v[:, :, len_t:]

        attn_t2t = (q_t @ k_t.transpose(-2, -1)) * self.scale
        attn_s2t = (q_s @ k_t.transpose(-2, -1)) * self.scale

        attn_t2t = attn_t2t.softmax(dim=-1)
        attn_s2t = attn_s2t.softmax(dim=-1)

        attn_t2t = self.attn_drop(attn_t2t)
        attn_s2t = self.attn_drop(attn_s2t)

        z = (attn_t2t @ v_t).transpose(1, 2).reshape(B, len_t, C)
        x = (attn_s2t @ v_t).transpose(1, 2).reshape(B, N-len_t, C)
        z = self.proj(z)
        x = self.proj(x)
        z = self.proj_drop(z)
        x = self.proj_drop(x)
        if return_attention:
            return torch.cat((z, x), dim=1), None
        else:
            return torch.cat((z, x), dim=1)

class CrossAttention(Attention):
    def forward(self, x, len_t=None, x_rel_pos_bias = None, return_attention = True):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q_t, k_t, v_t = q[:, :, :len_t], k[:, :, :len_t], v[:, :, :len_t]
        q_s, k_s, v_s = q[:, :, len_t:], k[:, :, len_t:], v[:, :, len_t:]

        attn_t2s = (q_t @ k_s.transpose(-2, -1)) * self.scale
        # attn_s2t = (q_s @ k_t.transpose(-2, -1)) * self.scale
        attn_s2t = attn_t2s.transpose(-1,-2)
        if x_rel_pos_bias is not None:
            # TODO Should be seprate RPB.
            attn_t2t = attn_t2t + x_rel_pos_bias.unsqueeze(0)
            attn_s2s = attn_s2s + x_rel_pos_bias.unsqueeze(0)

        attn_t2s = attn_t2s.softmax(dim=-1)
        attn_s2t = attn_s2t.softmax(dim=-1)
        attn_t2s = self.attn_drop(attn_t2s)
        attn_s2t = self.attn_drop(attn_s2t)

        z = (attn_t2s @ v_s).transpose(1, 2).reshape(B, len_t, C)
        x = (attn_s2t @ v_t).transpose(1, 2).reshape(B, N-len_t, C)
        z = self.proj(z)
        x = self.proj(x)
        z = self.proj_drop(z)
        x = self.proj_drop(x)
        if return_attention:
            return torch.cat((z, x), dim=1), None
        else:
            return torch.cat((z, x), dim=1)

class MixedCrossAttention(Attention):
    def forward(self, x, len_t, x_rel_pos_bias = None, return_attention = True):
        # z = x[:,:len_t]
        # x = x[:,len_t:]
        # len_t = len_t+1 # attend cls token
        # FIX: add cls outside
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q_s = q[:, :, len_t:]

        attn = (q_s @ k.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            attn = attn + x_rel_pos_bias.unsqueeze(0)

        attn = self.attn_drop(attn)
        # attn_t2t = attn[:, :, :len_t, :len_t]   # BS, Nh, len_t, len_t
        # attn_s2ts = attn[:, :, len_t:, :]       # BS, Nh, len_t+s, len_s
        attn_s2ts = attn

        # attn_t2t = attn_t2t.softmax(dim=-1)
        attn_s2ts = attn_s2ts.softmax(dim=-1)

        # x_t = (attn_t2t @ v[:, :, :len_t]).transpose(1, 2).reshape(B, len_t, C)
        x_t = x[:, :len_t]
        x_s = (attn_s2ts @ v).transpose(1, 2).reshape(B, N-len_t, C)
        x = torch.cat((x_t, x_s), dim=1)
        # x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        if return_attention:
            return x, attn
        else:
            return x

class HybridCrossAttention(Attention):
    # x: Concated tokens of:
    #       1. {POOLED tokens} from Template.
    #       2. {tokens} from Search.
    # Return: Search part of x.
    def forward(self, x, len_t=None, x_rel_pos_bias = None, return_attention = True, template_bb=None):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        # q_t, k_t, v_t = q[:, :, :len_t], k[:, :, :len_t], v[:, :, :len_t]
        q_s, k_s, v_s = q[:, :, len_t:], k[:, :, len_t:], v[:, :, len_t:]

        # attn_t2s = (q_t @ k_s.transpose(-2, -1)) * self.scale
        attn_s2ts = (q_s @ k.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            # TODO Should be seprate RPB.
            attn_t2t = attn_t2t + x_rel_pos_bias.unsqueeze(0)
            attn_s2s = attn_s2s + x_rel_pos_bias.unsqueeze(0)

        # attn_t2s = attn_t2s.softmax(dim=-1)
        attn_s2ts = attn_s2ts.softmax(dim=-1)
        # attn_t2s = self.attn_drop(attn_t2s)
        attn_s2ts = self.attn_drop(attn_s2ts)

        # z = (attn_t2s @ v_s).transpose(1, 2).reshape(B, len_t, C)
        x = (attn_s2ts @ v).transpose(1, 2).reshape(B, N-len_t, C)
        # z = self.proj(z)
        x = self.proj(x)
        # z = self.proj_drop(z)
        x = self.proj_drop(x)
        if return_attention:
            # return torch.cat((z, x), dim=1), None
            return x, None
        else:
            return x


class SeperateSelfAttention(Attention):
    def forward(self, x, len_t=None, x_rel_pos_bias = None, return_attention = True):
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) # 3, B, nH, N, dim
        q, k, v = qkv[0], qkv[1], qkv[2]

        q_t, k_t, v_t = q[:, :, :len_t], k[:, :, :len_t], v[:, :, :len_t]
        q_s, k_s, v_s = q[:, :, len_t:], k[:, :, len_t:], v[:, :, len_t:]

        attn_t2t = (q_t @ k_t.transpose(-2, -1)) * self.scale
        attn_s2s = (q_s @ k_s.transpose(-2, -1)) * self.scale
        if x_rel_pos_bias is not None:
            # TODO Should be seprate RPB.
            attn_t2t = attn_t2t + x_rel_pos_bias.unsqueeze(0)
            attn_s2s = attn_s2s + x_rel_pos_bias.unsqueeze(0)

        # attn_s2t = (q_s @ k_t.transpose(-2, -1)) * self.scale
        # attn_s2t = attn_s2t.softmax(dim=-1)
        # attn_s2t = self.attn_drop(attn_s2t)

        attn_t2t = attn_t2t.softmax(dim=-1)
        attn_s2s = attn_s2s.softmax(dim=-1)
        attn_t2t = self.attn_drop(attn_t2t)
        attn_s2s = self.attn_drop(attn_s2s)

        z = (attn_t2t @ v_t).transpose(1, 2).reshape(B, len_t, C)
        x = (attn_s2s @ v_s).transpose(1, 2).reshape(B, N-len_t, C)
        z = self.proj(z)
        x = self.proj(x)
        z = self.proj_drop(z)
        x = self.proj_drop(x)
        if return_attention:
            return torch.cat((z, x), dim=1), None
        else:
            return x

class CrossAttention_CAE(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
            proj_drop=0., window_size=None, attn_head_dim=None):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.q = nn.Linear(dim, all_head_dim, bias=False)
        self.k = nn.Linear(dim, all_head_dim, bias=False)
        self.v = nn.Linear(dim, all_head_dim, bias=False)

        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.k_bias = None
            self.v_bias = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(all_head_dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, bool_masked_pos=None, k=None, v=None):
        B, N, C = x.shape
        N_k = k.shape[1]
        N_v = v.shape[1]

        q_bias, k_bias, v_bias = None, None, None
        if self.q_bias is not None:
            q_bias = self.q_bias
            k_bias = torch.zeros_like(self.v_bias, requires_grad=False)
            v_bias = self.v_bias

        q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
        q = q.reshape(B, N, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)    # (B, N_head, N_q, dim)

        k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
        k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)

        v = F.linear(input=v, weight=self.v.weight, bias=v_bias)   
        v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))      # (B, N_head, N_q, N_k)
        
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, -1) 
        x = self.proj(x)
        x = self.proj_drop(x)

        return x

class AttentiveBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 window_size=None, attn_head_dim=None):
        super().__init__()

        self.norm1_q = norm_layer(dim)
        self.norm1_k = norm_layer(dim)
        self.norm1_v = norm_layer(dim)
        self.norm2_cross = norm_layer(dim)
        self.cross_attn =  CrossAttention_CAE(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        
    def forward(self, x_q, x_kv, pos_q, pos_k, bool_masked_pos, rel_pos_bias=None):
        x_q = self.norm1_q(x_q + pos_q)
        x_k = self.norm1_k(x_kv + pos_k)
        x_v = self.norm1_v(x_kv)

        x = self.cross_attn(x_q, k=x_k, v=x_v)

        return x

class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., seperate_self_attn=False,
                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, window_size=None, init_values=0,
                 add_cls_token=False):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.add_cls_token = add_cls_token
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.seperate_self_attn = seperate_self_attn
        if not self.seperate_self_attn:
            self.attn = MixedCrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
                    attn_drop=attn_drop, proj_drop=drop, window_size=window_size)
        else:
            self.attn = HybridCrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
                    attn_drop=attn_drop, proj_drop=drop, window_size=window_size)
        if init_values > 0:
            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
        else:
            self.gamma_1, self.gamma_2 = None, None

    def forward(self, x, len_t, x_rel_pos_bias=None, return_attention=False, mask=None):
        # if self.seperate_self_attn:
        #     t = x[:, :len_t]
        #     s = x[:, len_t:]
        #     t_out, t_attn = self.attn(self.norm1(t), len_t, x_rel_pos_bias) # If ori Attn, 'len_t' is not used.
        #     s_out, s_attn = self.attn(self.norm1(s), len_t, x_rel_pos_bias)
        #     y = torch.cat((t_out, s_out), dim=1)
        #     attn = s_attn
        # else:
        z_in = x[:, :len_t]
        x_in = x[:, len_t:]
        # in_attn = torch.cat((z_in, self.norm1(x[:, len_t:])), dim=1)
        
        y, attn = self.attn(self.norm1(x), len_t, x_rel_pos_bias)
        y_out = y[:, len_t:]
        # x = x[:, len_t:]
        if return_attention:
            return attn
        if self.gamma_1 is None:
            x = x + self.drop_path(y)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x_in = x_in + self.drop_path(self.gamma_1 * y_out)
            x_in = x_in + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x_in)))
        out = torch.cat((z_in, x_in), dim=1)
        return out, attn

class CAE_CEBlock(Block):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., keep_ratio_search=1.0,
                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, window_size=None, init_values=0):
        super().__init__(dim = dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop,
        attn_drop=attn_drop, drop_path=drop_path, act_layer=act_layer, norm_layer=norm_layer, window_size=window_size, init_values=init_values)
        self.keep_ratio_search = keep_ratio_search
    def forward(self, x, global_index_template, global_index_search, x_rel_pos_bias=None, mask=None, ce_template_mask=None, keep_ratio_search=None, use_target_token=False):
        lens_t = global_index_template.shape[1]
        lens_all = x.shape[1]
        # if use_target_token:
        #     mask = generate_mask_for_target_token(mask, lens_all, lens_t)
        # else:
        #     mask = None
        y, attn = self.attn(self.norm1(x), x_rel_pos_bias)
        if self.gamma_1 is None:
            x = x + self.drop_path(y)
        else:
            x = x + self.drop_path(self.gamma_1 * y)
        # if use_target_token:
        #     target_token = x[:,-1:]
        #     x = x[:,:-1]

        removed_index_search = None
        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)
        # x = x + self.drop_path(self.mlp(self.norm2(x)))
        if self.gamma_1 is None:
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        # if use_target_token:
        #     x = torch.cat([x, target_token], dim=1)
        return x, global_index_template, global_index_search, removed_index_search, attn

def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor, lens_t: int, keep_ratio: float, global_index: torch.Tensor, box_mask_z: torch.Tensor):
    """
    Eliminate potential background candidates for computation reduction and noise cancellation.
    Args:
        attn (torch.Tensor): [B, num_heads, L_t + L_s, L_t + L_s], attention weights
        tokens (torch.Tensor):  [B, L_t + L_s, C], template and search region tokens
        lens_t (int): length of template
        keep_ratio (float): keep ratio of search region tokens (candidates)
        global_index (torch.Tensor): global index of search region tokens
        box_mask_z (torch.Tensor): template mask used to accumulate attention weights

    Returns:
        tokens_new (torch.Tensor): tokens after candidate elimination
        keep_index (torch.Tensor): indices of kept search region tokens
        removed_index (torch.Tensor): indices of removed search region tokens
    """
    lens_s = attn.shape[-1] - lens_t
    bs, hn, _, _ = attn.shape

    lens_keep = math.ceil(keep_ratio * lens_s)
    if lens_keep == lens_s:
        return tokens, global_index, None

    attn_t = attn[:, :, :lens_t, lens_t:]

    if box_mask_z is not None:
        box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(-1, attn_t.shape[1], -1, attn_t.shape[-1])
        # attn_t = attn_t[:, :, box_mask_z, :]
        attn_t = attn_t[box_mask_z]
        attn_t = attn_t.view(bs, hn, -1, lens_s)
        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s

        # attn_t = [attn_t[i, :, box_mask_z[i, :], :] for i in range(attn_t.size(0))]
        # attn_t = [attn_t[i].mean(dim=1).mean(dim=0) for i in range(len(attn_t))]
        # attn_t = torch.stack(attn_t, dim=0)
    else:
        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s

    # use sort instead of topk, due to the speed issue
    # https://github.com/pytorch/pytorch/issues/22812
    sorted_attn, indices = torch.sort(attn_t, dim=1, descending=True)

    topk_attn, topk_idx = sorted_attn[:, :lens_keep], indices[:, :lens_keep]
    non_topk_attn, non_topk_idx = sorted_attn[:, lens_keep:], indices[:, lens_keep:]

    keep_index = global_index.gather(dim=1, index=topk_idx)
    removed_index = global_index.gather(dim=1, index=non_topk_idx)

    # separate template and search tokens
    tokens_t = tokens[:, :lens_t]
    tokens_s = tokens[:, lens_t:]

    # obtain the attentive and inattentive tokens
    B, L, C = tokens_s.shape
    # topk_idx_ = topk_idx.unsqueeze(-1).expand(B, lens_keep, C)
    attentive_tokens = tokens_s.gather(dim=1, index=topk_idx.unsqueeze(-1).expand(B, -1, C))
    # inattentive_tokens = tokens_s.gather(dim=1, index=non_topk_idx.unsqueeze(-1).expand(B, -1, C))

    # compute the weighted combination of inattentive tokens
    # fused_token = non_topk_attn @ inattentive_tokens

    # concatenate these tokens
    # tokens_new = torch.cat([tokens_t, attentive_tokens, fused_token], dim=0)
    tokens_new = torch.cat([tokens_t, attentive_tokens], dim=1)

    return tokens_new, keep_index, removed_index

def generate_mask_for_target_token(center_attn_mask, len_all, len_t):

    # len_inner = len_t + 1
    center_attn_mask = center_attn_mask[:len_all,:len_all]
    center_attn_mask[:,:] = True
    # center_attn_mask = torch.ones([len_all, len_all], dtype=torch.bool)
    # center_attn_mask[:len_inner, :len_inner] = True # Self attn in Template and TargetToken
    # center_attn_mask[len_inner:, len_inner:] = True # Self attn in Search
    center_attn_mask[len_t:-1, -1:] = False #  
    center_attn_mask[-1:, len_t:-1] = False #
    return center_attn_mask.detach()

class CEBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, keep_ratio_search=1.0,):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

        self.keep_ratio_search = keep_ratio_search

    def forward(self, x, global_index_template, global_index_search, mask=None, ce_template_mask=None, keep_ratio_search=None):
        x_attn, attn = self.attn(self.norm1(x), mask, True)
        x = x + self.drop_path(x_attn)
        lens_t = global_index_template.shape[1]

        removed_index_search = None
        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)

        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x, global_index_template, global_index_search, removed_index_search, attn

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=[224, 224], patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.num_patches_w = img_size[0] // patch_size
        self.num_patches_h = img_size[1] // patch_size

        num_patches = self.num_patches_w * self.num_patches_h
        self.patch_shape = (img_size[0] // patch_size, img_size[1] // patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        # print("##############patch here!!!")
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.flatten = True
            
    def forward(self, x, mask=None):
        # B, C, H, W = x.shape
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        # x = self.norm(x)
        return x


class VIT_Backbone(nn.Module):
    def __init__(self,
                 img_size=[224,224],
                 patch_size=16,
                 embed_dim=768,
                 in_chans=3,
                 frozen_stages=0,
                #  out_indices=[3, 5, 7, 11],
                 out_with_norm=False,
                 use_checkpoint=False,
                 CE=False,
                 add_cls_token = False,
                 add_template_target_embed=False,
                 num_classes=0,  
                 depth=12,
                 num_heads=12, 
                 mlp_ratio=4., 
                 qkv_bias=False, 
                 qk_scale=None, 
                 drop_rate=0., 
                 attn_drop_rate=0.,
                 drop_path_rate=0., 
                 norm_layer=partial(nn.LayerNorm, eps=1e-6), 
                 return_all_tokens=False, 
                 init_values=0, 
                 use_sincos_pos_emb=False, 
                 use_abs_pos_emb=False, 
                 use_rel_pos_bias=False, 
                 seperate_loc=None,
                 act_layer=None, 
                 ce_loc=None, 
                 ce_keep_ratio=None, 
                 use_mean_pooling=False, 
                 masked_im_modeling=False,
                 ):
        
        super().__init__()

        self.num_features = self.embed_dim = embed_dim
        self.return_all_tokens = return_all_tokens
        # print("############use_abs_pos:", use_abs_pos_emb)
        # print("############use_sincos_pos:", use_sincos_pos_emb)
        # print("############use_rel_pos_bias:", use_rel_pos_bias)
        self.use_abs_pos_emb = use_abs_pos_emb
        self.use_sincos_pos_emb = use_sincos_pos_emb
        self.CE = CE
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        if use_abs_pos_emb:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        else:
            self.pos_embed = None

        self.pos_drop = nn.Dropout(p=drop_rate)
        act_layer = act_layer or nn.GELU
        self.use_rel_pos_bias = use_rel_pos_bias

        if self.use_rel_pos_bias:
            print("=================use RelativePositionBias===================")
            window_size=self.patch_embed.patch_shape
            self.window_size = window_size
            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
            self.relative_position_bias_table = nn.Parameter(
                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
            # cls to token & token 2 cls & cls to cls

        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        if self.CE:    
            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
            blocks = []
            ce_index = 0
            self.ce_loc = ce_loc
            for i in range(depth):
                ce_keep_ratio_i = 1.0
                if ce_loc is not None and i in ce_loc:
                    ce_keep_ratio_i = ce_keep_ratio[ce_index]
                    ce_index += 1

                blocks.append(
                    CAE_CEBlock(
                        dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                        drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 
                        init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
                        keep_ratio_search=ce_keep_ratio_i)
                )
            self.blocks = nn.ModuleList(blocks)
        else:
            blocks = []
            sep_index = 0
            self.seperate_loc = seperate_loc
            for i in range(depth):
                sep_i = False
                if seperate_loc is not None and i in seperate_loc:
                    sep_i = True
                    sep_index += 1
                blocks.append(Block(
                    dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                    drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, seperate_self_attn = sep_i,
                    init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None))
            self.blocks = nn.ModuleList(blocks)

        # self.attentive_blocks_ = nn.ModuleList([
        #             AttentiveBlock(
        #                 dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
        #                 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=0, norm_layer=norm_layer,
        #                 init_values=0)
        #             for i in range(3)])    

        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
        # Classifier head
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
        
        if use_abs_pos_emb:
            trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        self.apply(self._init_weights)

        # masked image modeling
        self.masked_im_modeling = masked_im_modeling
        if masked_im_modeling:
            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))

        # support non-square image as input
        if len(img_size) == 1:
            img_size = img_size * 2
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches
        if self.use_abs_pos_emb:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        elif self.use_sincos_pos_emb:
            self.pos_embed = self.build_2d_sincos_position_embedding(embed_dim)
        else:
            self.pos_embed = None

        self.add_cls_token = add_cls_token
        self.patch_size = patch_size
        # self.with_fpn = with_fpn
        self.frozen_stages = frozen_stages
        # self.out_indices = out_indices
        self.use_checkpoint = use_checkpoint
        self.cat_mode = 'direct'
        self.add_template_target_embed = add_template_target_embed

        self.pos_embed_z = None
        self.pos_embed_x = None
        # if self.add_cls_token:
        #     # assert self.CE == False # Not implement yet for use both.
        #     # self.cls_embed = Mlp(4, out_features=self.embed_dim)

        if not out_with_norm:
            self.norm = nn.Identity()
        # self.prpool = PrRoIPool([2,2])

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def prepare_tokens(self, x, mask=None):
        B, nc, w, h = x.shape
        # patch linear embedding
        x = self.patch_embed(x)

        # mask image modeling
        if mask is not None:
            x = self.mask_model(x, mask)
        x = x.flatten(2).transpose(1, 2)

        # add the [CLS] token to the embed patch tokens
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        # add positional encoding to each token
        if self.pos_embed is not None:
            x = x + self.interpolate_pos_encoding(x, w, h)

        return self.pos_drop(x)

    def get_last_selfattention(self, x):
        x = self.prepare_tokens(x)
        for i, blk in enumerate(self.blocks):
            if i < len(self.blocks) - 1:
                x = blk(x)
            else:
                # return attention of the last block
                return blk(x, return_attention=True)

    def get_intermediate_layers(self, x, n=1):
        x = self.prepare_tokens(x)
        # we return the output tokens from the `n` last blocks
        output = []
        for i, blk in enumerate(self.blocks):
            x = blk(x)
            if len(self.blocks) - i <= n:
                output.append(self.norm(x))
        return output
        
    def get_num_layers(self):
        return len(self.blocks)

    def build_2d_sincos_position_embedding(self, embed_dim=768, temperature=10000., decode=False):
        h, w = self.patch_embed.patch_shape 
        grid_w = torch.arange(w, dtype=torch.float32)
        grid_h = torch.arange(h, dtype=torch.float32)
        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
        assert embed_dim % 4 == 0, 'Embed dimension must be divisible by 4 for 2D sin-cos position embedding'
        pos_dim = embed_dim // 4
        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
        omega = 1. / (temperature ** omega)
        out_w = torch.einsum('m,d->md', [grid_w.flatten(), omega])
        out_h = torch.einsum('m,d->md', [grid_h.flatten(), omega])
        pos_emb = torch.cat([torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)], dim=1)[None, :, :]

        pe_token = torch.zeros([1, 1, embed_dim], dtype=torch.float32)
        pos_embed = nn.Parameter(torch.cat([pe_token, pos_emb], dim=1))
        pos_embed.requires_grad = False
        return pos_embed
    
    def finetune_track(self, cfg, patch_start_index=1):
        # patch_start_index: For skipping the cls embedding
        search_size = to_2tuple(cfg.DATA.SEARCH.SIZE)
        template_size = to_2tuple(cfg.DATA.TEMPLATE.SIZE)
        new_patch_size = cfg.MODEL.BACKBONE.STRIDE

        self.cat_mode = cfg.MODEL.BACKBONE.CAT_MODE
        self.return_inter = cfg.MODEL.RETURN_INTER
        self.add_sep_seg = cfg.MODEL.BACKBONE.SEP_SEG

        # resize patch embedding
        if new_patch_size != self.patch_size:
            print('Inconsistent Patch Size With The Pretrained Weights, Interpolate The Weight!')
            old_patch_embed = {}
            for name, param in self.patch_embed.named_parameters():
                if 'weight' in name:
                    param = nn.functional.interpolate(param, size=(new_patch_size, new_patch_size),
                                                      mode='bicubic', align_corners=False)
                    param = nn.Parameter(param)
                old_patch_embed[name] = param
            # self.patch_embed = PatchEmbed(img_size=self.img_size, patch_size=new_patch_size, in_chans=3,
            #                               embed_dim=self.embed_dim)
            self.patch_embed.proj.bias = old_patch_embed['proj.bias']
            self.patch_embed.proj.weight = old_patch_embed['proj.weight']

        # Pos embedding
        patch_pos_embed = self.pos_embed[:, patch_start_index:, :]
        patch_pos_embed = patch_pos_embed.transpose(1, 2) #?
        B, E, Q = patch_pos_embed.shape
        P_H, P_W = self.patch_embed.patch_shape[0], self.patch_embed.patch_shape[1]
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)

        # Pos embedding for search region
        H, W = search_size
        new_P_H_S, new_P_W_S = H // new_patch_size, W // new_patch_size
        search_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H_S, new_P_W_S), mode='bicubic',
                                                           align_corners=False)
        search_patch_pos_embed = search_patch_pos_embed.flatten(2).transpose(1, 2)

        # Pos embedding for template region
        H, W = template_size
        new_P_H_T, new_P_W_T = H // new_patch_size, W // new_patch_size
        template_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H_T, new_P_W_T), mode='bicubic',
                                                             align_corners=False)
        template_patch_pos_embed = template_patch_pos_embed.flatten(2).transpose(1, 2)

        self.pos_embed_z = nn.Parameter(template_patch_pos_embed)
        self.pos_embed_x = nn.Parameter(search_patch_pos_embed)

        # for cls token 
        if self.add_cls_token and patch_start_index > 0:
            cls_pos_embed = self.pos_embed[:, 0:1, :]
            self.cls_pos_embed = nn.Parameter(cls_pos_embed)

        # separate token and segment token
        if self.add_sep_seg:
            self.template_segment_pos_embed = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
            self.template_segment_pos_embed = trunc_normal_(self.template_segment_pos_embed, std=.02)
            self.search_segment_pos_embed = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
            self.search_segment_pos_embed = trunc_normal_(self.search_segment_pos_embed, std=.02)

        # self.cls_token = None
        # self.pos_embed = None

        if self.return_inter:
            for i_layer in self.fpn_stage:
                if i_layer != 11:
                    norm_layer = partial(nn.LayerNorm, eps=1e-6)
                    layer = norm_layer(self.embed_dim)
                    layer_name = f'norm{i_layer}'
                    self.add_module(layer_name, layer)

        if self.add_template_target_embed:
            self.template_target_embed = PatchEmbed(
                img_size=template_size, patch_size=new_patch_size, in_chans=1)

        # For 8x8 tokens in Template and 16x16 tokens in Search, the size of attn map is 8x8+16x16 = 320
        attn_map_size = new_P_W_T * new_P_H_T + new_P_W_S * new_P_H_S   
        if self.add_cls_token:
            attn_map_size += 1
        center_attn_mask = torch.ones([attn_map_size, attn_map_size], dtype=torch.bool)
        # center_attn_mask[:new_P_W_T * new_P_H_T, :new_P_W_T * new_P_H_T] = True # Self attn in Template
        # center_attn_mask[new_P_W_T * new_P_H_T:, new_P_W_T * new_P_H_T:] = True # Self attn in Search
        # center_attn_mask[new_P_W_T * new_P_H_T//2] = True # Template to Search attn 
        # center_attn_mask[:, new_P_W_T * new_P_H_T//2] = True # Search to Template attn
        self.register_buffer('center_attn_mask', center_attn_mask)
   
    def train(self, mode=True):
        """Convert the model into training mode while keep layers freezed."""
        super(VIT_Backbone, self).train(mode)
        self._freeze_stages()
        # if self.pos_embed is not None:
        #     if self.pos_embed.requires_grad:
        #         print("=================pos_embed update ================")
        #     else:
        #         print("=================pos_embed static ================")
            
    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.patch_embed.eval()
            for param in self.patch_embed.parameters():
                param.requires_grad = False
            self.cls_token.requires_grad = False
            if self.pos_embed is not None and self.use_sincos_pos_emb == True:
                self.pos_embed.requires_grad = False
            self.pos_drop.eval()

        for i in range(1, self.frozen_stages + 1):
            
            if i  == len(self.blocks):
                norm_layer = getattr(self, 'norm') #f'norm{i-1}')
                norm_layer.eval()
                for param in norm_layer.parameters():
                    param.requires_grad = False

            m = self.blocks[i - 1]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False
        
    def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        if npatch == N and w0 == self.patch_embed.num_patches_w and h0 == self.patch_embed.num_patches_h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        
        tmp=patch_pos_embed.reshape(1, self.patch_embed.num_patches_w, self.patch_embed.num_patches_h, dim).permute(0, 3, 1, 2)
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, self.patch_embed.num_patches_w, self.patch_embed.num_patches_h, dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / self.patch_embed.num_patches_w, h0 / self.patch_embed.num_patches_h),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        # return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
        return patch_pos_embed

    def forward_features(self, z, x, mask_z=None, mask_x=None,
                         ce_template_mask=None, ce_keep_rate=None,
                         return_last_attn=False, template_bb=None,
                         target_token=None
                         ):
        B, H, W = x.shape[0], x.shape[2], x.shape[3]

        x = self.patch_embed(x)
        # z = self.patch_embed(z)
        if self.add_template_target_embed:
            target_embed = self.template_target_embed(mask_z.unsqueeze(1).float())
            z += target_embed

        # attention mask handling
        # B, H, W
        # if mask_z is not None and mask_x is not None:
        #     mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
        #     mask_z = mask_z.flatten(1).unsqueeze(-1)

        #     mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
        #     mask_x = mask_x.flatten(1).unsqueeze(-1)

        #     mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode)
        #     mask_x = mask_x.squeeze(-1)

        # if self.add_cls_token:
        #     cls_tokens = self.cls_embed(template_bb).unsqueeze(-2)
            # cls_tokens = self.cls_token.expand(B, -1, -1)
            # cls_tokens = cls_tokens + self.cls_pos_embed

        # z += self.pos_embed_z
        x += self.pos_embed_x

        if self.add_sep_seg:
            x += self.search_segment_pos_embed
            z += self.template_segment_pos_embed

        if self.add_cls_token:
            z = torch.cat([target_token, z], dim=1)
        # x = combine_tokens(z, x, mode=self.cat_mode)       
        
        x = self.pos_drop(x)
        # z = self.pos_drop(z)

        lens_z = self.pos_embed_z.shape[1]
        lens_x = self.pos_embed_x.shape[1]

        # global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device)
        # global_index_t = global_index_t.repeat(B, 1)

        # global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device)
        # global_index_s = global_index_s.repeat(B, 1)
        removed_indexes_s = []
        attns = []
        if self.CE:
            removed_indexes_s = []
            for i, blk in enumerate(self.blocks):
                x, global_index_t, global_index_s, removed_index_s, attn = \
                    blk(x, global_index_t, global_index_s, x_rel_pos_bias=None,
                        mask=self.center_attn_mask, ce_template_mask=ce_template_mask,
                        keep_ratio_search=ce_keep_rate, use_target_token=self.add_cls_token)

                if self.ce_loc is not None and i in self.ce_loc:
                    removed_indexes_s.append(removed_index_s)
        else:
            # z = x[:, :lens_z].to(torch.float32)
            # x = x[:, lens_z:]

            # # Async:
            # for i in range(0,9):
            #     z, attn = self.blocks[i](z, 0, None,  mask=None)
            if self.add_cls_token:
                z = z[:,1:]
            for i in range(0,4):
                x, attn = self.blocks[i](x, 0, None,  mask=None)
            x = combine_tokens(z, x, mode=self.cat_mode)
            for i in range(4,6):
                x, attn = self.blocks[i](x, lens_z, None,  mask=None)
            
            # One-Stream:
            # for i in range(0,12):
            #     x, attn = self.blocks[i](x, 0, None,  mask=None)

            

            # x = torch.cat([z, x], dim=1)

            # # attentive_index = 0
            # # for i, blk in enumerate(self.blocks):
            # #     if i in [5, 7, 11]:
            # x = self.attentive_blocks[0](x,z,0,0, bool_masked_pos=None, rel_pos_bias=None)
            #         attentive_index += 1
            #     x, attn = blk(x, lens_z, None,  mask=self.center_attn_mask)
        x = self.norm(x)
        # lens_x_new = global_index_s.shape[1]
        # lens_z_new = global_index_t.shape[1]
        if self.add_cls_token:
            z = x[:, :lens_z]
            x = x[:, lens_z:]
        else:
            z = x[:, :lens_z]
            x = x[:, lens_z:]

        if removed_indexes_s and removed_indexes_s[0] is not None:
            removed_indexes_cat = torch.cat(removed_indexes_s, dim=1)

            pruned_lens_x = lens_x - lens_x_new
            pad_x = torch.zeros([B, pruned_lens_x, x.shape[2]], device=x.device)
            x = torch.cat([x, pad_x], dim=1)
            index_all = torch.cat([global_index_s, removed_indexes_cat], dim=1)
            # recover original token order
            C = x.shape[-1]
            # x = x.gather(1, index_all.unsqueeze(-1).expand(B, -1, C).argsort(1))
            x = torch.zeros_like(x).scatter_(dim=1, index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64), src=x)

        # x = recover_tokens(x, lens_z_new, lens_x, mode=self.cat_mode)

        # re-concatenate with the template, which may be further used by other modules
        x = torch.cat([z, x], dim=1)
        # attns = torch.cat(attns, dim=1)
        aux_dict = {
            # "attn": attns[:,:,lens_z:,:lens_z].mean(-1).reshape(B,-1,16,16),
            "attn": attn,
            "removed_indexes_s": removed_indexes_s,  # used for visualization
        }

        return x, aux_dict

    def forward(self, z, x, ce_template_mask=None, ce_keep_rate=None,
                tnc_keep_rate=None,
                return_last_attn=False,
                template_bb=None,
                mask_z=None,
                target_token=None):

        x, aux_dict = self.forward_features(z, x, ce_template_mask=ce_template_mask,
                ce_keep_rate=ce_keep_rate,template_bb=template_bb, mask_z=mask_z, target_token=target_token)

        return x, aux_dict



def _create_ViT_CAE(pretrained=False, **kwargs) -> VIT_Backbone: 
    model = VIT_Backbone(**kwargs)

    if pretrained:
        if 'npz' in pretrained:
            model.load_pretrained(pretrained, prefix='')
        else:
            state_dict = torch.load('pretrained_models/cae_base.pth', map_location="cpu")
            # new_ckpt = OrderedDict()
            state_dict = state_dict['model']
            if sorted(list(state_dict.keys()))[0].startswith('encoder'):
                state_dict = {k.replace('encoder.', ''): v for k, v in state_dict.items() if k.startswith('encoder.')}
            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
            print("Unexpected keys:")
            print(unexpected_keys)
            print("Missing keys:")
            print(missing_keys)
    return model

def CAE_Base_patch16_224(pretrained=False, **kwargs) -> VIT_Backbone:
    model_kwargs = dict(mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
        patch_size=16, embed_dim=768, depth=12, num_heads=12,
        drop_rate=0., attn_drop_rate=0., init_values=0.1, 
        use_abs_pos_emb=True, **kwargs)
    model = _create_ViT_CAE(pretrained=pretrained, **model_kwargs)
    # model.adpat4track()
    return model
lib/models/ostrack/vit_cae.py
TsingWei/LiteTrack
import math
import logging
from functools import partial
from collections import OrderedDict
from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F

from timm.models.layers import to_2tuple

from lib.models.layers.patch_embed import PatchEmbed
from .utils import combine_tokens, recover_tokens
from .vit import VisionTransformer
from ..layers.attn_blocks import CEBlock

_logger = logging.getLogger(__name__)


class VisionTransformerCE(VisionTransformer):
    """ Vision Transformer with candidate elimination (CE) module

    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
        - https://arxiv.org/abs/2010.11929

    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
        - https://arxiv.org/abs/2012.12877
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
                 act_layer=None, weight_init='',
                 ce_loc=None, ce_keep_ratio=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
            weight_init: (str): weight init scheme
        """
        # super().__init__()
        super().__init__()
        if isinstance(img_size, tuple):
            self.img_size = img_size
        else:
            self.img_size = to_2tuple(img_size)
        self.patch_size = patch_size
        self.in_chans = in_chans

        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_tokens = 2 if distilled else 1
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        act_layer = act_layer or nn.GELU

        self.patch_embed = embed_layer(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        blocks = []
        ce_index = 0
        self.ce_loc = ce_loc
        for i in range(depth):
            ce_keep_ratio_i = 1.0
            if ce_loc is not None and i in ce_loc:
                ce_keep_ratio_i = ce_keep_ratio[ce_index]
                ce_index += 1

            blocks.append(
                CEBlock(
                    dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
                    attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer,
                    keep_ratio_search=ce_keep_ratio_i)
            )

        self.blocks = nn.Sequential(*blocks)
        self.norm = norm_layer(embed_dim)

        self.init_weights(weight_init)

    def forward_features(self, z, x, mask_z=None, mask_x=None,
                         ce_template_mask=None, ce_keep_rate=None,
                         return_last_attn=False
                         ):
        B, H, W = x.shape[0], x.shape[2], x.shape[3]

        x = self.patch_embed(x)
        z = self.patch_embed(z)

        # attention mask handling
        # B, H, W
        if mask_z is not None and mask_x is not None:
            mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
            mask_z = mask_z.flatten(1).unsqueeze(-1)

            mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
            mask_x = mask_x.flatten(1).unsqueeze(-1)

            mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode)
            mask_x = mask_x.squeeze(-1)

        if self.add_cls_token:
            cls_tokens = self.cls_token.expand(B, -1, -1)
            cls_tokens = cls_tokens + self.cls_pos_embed

        z += self.pos_embed_z
        x += self.pos_embed_x

        if self.add_sep_seg:
            x += self.search_segment_pos_embed
            z += self.template_segment_pos_embed

        x = combine_tokens(z, x, mode=self.cat_mode)
        if self.add_cls_token:
            x = torch.cat([cls_tokens, x], dim=1)

        x = self.pos_drop(x)

        lens_z = self.pos_embed_z.shape[1]
        lens_x = self.pos_embed_x.shape[1]

        global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device)
        global_index_t = global_index_t.repeat(B, 1)

        global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device)
        global_index_s = global_index_s.repeat(B, 1)
        removed_indexes_s = []
        for i, blk in enumerate(self.blocks):
            x, global_index_t, global_index_s, removed_index_s, attn = \
                blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate)

            if self.ce_loc is not None and i in self.ce_loc:
                removed_indexes_s.append(removed_index_s)

        x = self.norm(x)
        lens_x_new = global_index_s.shape[1]
        lens_z_new = global_index_t.shape[1]

        z = x[:, :lens_z_new]
        x = x[:, lens_z_new:]

        if removed_indexes_s and removed_indexes_s[0] is not None:
            removed_indexes_cat = torch.cat(removed_indexes_s, dim=1)

            pruned_lens_x = lens_x - lens_x_new
            pad_x = torch.zeros([B, pruned_lens_x, x.shape[2]], device=x.device)
            x = torch.cat([x, pad_x], dim=1)
            index_all = torch.cat([global_index_s, removed_indexes_cat], dim=1)
            # recover original token order
            C = x.shape[-1]
            # x = x.gather(1, index_all.unsqueeze(-1).expand(B, -1, C).argsort(1))
            x = torch.zeros_like(x).scatter_(dim=1, index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64), src=x)

        x = recover_tokens(x, lens_z_new, lens_x, mode=self.cat_mode)

        # re-concatenate with the template, which may be further used by other modules
        x = torch.cat([z, x], dim=1)

        aux_dict = {
            "attn": attn,
            "removed_indexes_s": removed_indexes_s,  # used for visualization
        }

        return x, aux_dict

    def forward(self, z, x, ce_template_mask=None, ce_keep_rate=None,
                tnc_keep_rate=None,
                return_last_attn=False):

        x, aux_dict = self.forward_features(z, x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,)

        return x, aux_dict


def _create_vision_transformer(pretrained=False, **kwargs) -> VisionTransformerCE: 
    model = VisionTransformerCE(**kwargs)

    if pretrained:
        if 'npz' in pretrained:
            model.load_pretrained(pretrained, prefix='')
        else:
            checkpoint = torch.load(pretrained, map_location="cpu")
            missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False)
            print('Load pretrained model from: ' + pretrained)

    return model


def vit_base_patch16_224_ce(pretrained=False, **kwargs) -> VisionTransformerCE:
    """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    """
    model_kwargs = dict(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
    return model


def vit_large_patch16_224_ce(pretrained=False, **kwargs) -> VisionTransformerCE:
    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
    """
    model_kwargs = dict(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)
    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
    return model

lib/models/ostrack/vit_ce.py
TsingWei/LiteTrack
import tikzplotlib
import matplotlib
import matplotlib.pyplot as plt
import os
import torch
import pickle
import json
from lib.test.evaluation.environment import env_settings
from lib.test.analysis.extract_results import extract_results


def get_plot_draw_styles():
    plot_draw_style = [{'color': (1.0, 0.0, 0.0), 'line_style': '-'},
                       {'color': (0.0, 1.0, 0.0), 'line_style': '-'},
                       {'color': (0.0, 0.0, 1.0), 'line_style': '-'},
                       {'color': (0.0, 0.0, 0.0), 'line_style': '-'},
                       {'color': (1.0, 0.0, 1.0), 'line_style': '-'},
                       {'color': (0.0, 1.0, 1.0), 'line_style': '-'},
                       {'color': (0.5, 0.5, 0.5), 'line_style': '-'},
                       {'color': (136.0 / 255.0, 0.0, 21.0 / 255.0), 'line_style': '-'},
                       {'color': (1.0, 127.0 / 255.0, 39.0 / 255.0), 'line_style': '-'},
                       {'color': (0.0, 162.0 / 255.0, 232.0 / 255.0), 'line_style': '-'},
                       {'color': (0.0, 0.5, 0.0), 'line_style': '-'},
                       {'color': (1.0, 0.5, 0.2), 'line_style': '-'},
                       {'color': (0.1, 0.4, 0.0), 'line_style': '-'},
                       {'color': (0.6, 0.3, 0.9), 'line_style': '-'},
                       {'color': (0.4, 0.7, 0.1), 'line_style': '-'},
                       {'color': (0.2, 0.1, 0.7), 'line_style': '-'},
                       {'color': (0.7, 0.6, 0.2), 'line_style': '-'}]

    return plot_draw_style


def check_eval_data_is_valid(eval_data, trackers, dataset):
    """ Checks if the pre-computed results are valid"""
    seq_names = [s.name for s in dataset]
    seq_names_saved = eval_data['sequences']

    tracker_names_f = [(t.name, t.parameter_name, t.run_id) for t in trackers]
    tracker_names_f_saved = [(t['name'], t['param'], t['run_id']) for t in eval_data['trackers']]

    return seq_names == seq_names_saved and tracker_names_f == tracker_names_f_saved


def merge_multiple_runs(eval_data):
    new_tracker_names = []
    ave_success_rate_plot_overlap_merged = []
    ave_success_rate_plot_center_merged = []
    ave_success_rate_plot_center_norm_merged = []
    avg_overlap_all_merged = []

    ave_success_rate_plot_overlap = torch.tensor(eval_data['ave_success_rate_plot_overlap'])
    ave_success_rate_plot_center = torch.tensor(eval_data['ave_success_rate_plot_center'])
    ave_success_rate_plot_center_norm = torch.tensor(eval_data['ave_success_rate_plot_center_norm'])
    avg_overlap_all = torch.tensor(eval_data['avg_overlap_all'])

    trackers = eval_data['trackers']
    merged = torch.zeros(len(trackers), dtype=torch.uint8)
    for i in range(len(trackers)):
        if merged[i]:
            continue
        base_tracker = trackers[i]
        new_tracker_names.append(base_tracker)

        match = [t['name'] == base_tracker['name'] and t['param'] == base_tracker['param'] for t in trackers]
        match = torch.tensor(match)

        ave_success_rate_plot_overlap_merged.append(ave_success_rate_plot_overlap[:, match, :].mean(1))
        ave_success_rate_plot_center_merged.append(ave_success_rate_plot_center[:, match, :].mean(1))
        ave_success_rate_plot_center_norm_merged.append(ave_success_rate_plot_center_norm[:, match, :].mean(1))
        avg_overlap_all_merged.append(avg_overlap_all[:, match].mean(1))

        merged[match] = 1

    ave_success_rate_plot_overlap_merged = torch.stack(ave_success_rate_plot_overlap_merged, dim=1)
    ave_success_rate_plot_center_merged = torch.stack(ave_success_rate_plot_center_merged, dim=1)
    ave_success_rate_plot_center_norm_merged = torch.stack(ave_success_rate_plot_center_norm_merged, dim=1)
    avg_overlap_all_merged = torch.stack(avg_overlap_all_merged, dim=1)

    eval_data['trackers'] = new_tracker_names
    eval_data['ave_success_rate_plot_overlap'] = ave_success_rate_plot_overlap_merged.tolist()
    eval_data['ave_success_rate_plot_center'] = ave_success_rate_plot_center_merged.tolist()
    eval_data['ave_success_rate_plot_center_norm'] = ave_success_rate_plot_center_norm_merged.tolist()
    eval_data['avg_overlap_all'] = avg_overlap_all_merged.tolist()

    return eval_data


def get_tracker_display_name(tracker):
    if tracker['disp_name'] is None:
        if tracker['run_id'] is None:
            disp_name = '{}_{}'.format(tracker['name'], tracker['param'])
        else:
            disp_name = '{}_{}_{:03d}'.format(tracker['name'], tracker['param'],
                                              tracker['run_id'])
    else:
        disp_name = tracker['disp_name']

    return  disp_name


def plot_draw_save(y, x, scores, trackers, plot_draw_styles, result_plot_path, plot_opts):
    plt.rcParams['text.usetex']=True
    plt.rcParams["font.family"] = "Times New Roman"
    # Plot settings
    font_size = plot_opts.get('font_size', 20)
    font_size_axis = plot_opts.get('font_size_axis', 20)
    line_width = plot_opts.get('line_width', 2)
    font_size_legend = plot_opts.get('font_size_legend', 20)

    plot_type = plot_opts['plot_type']
    legend_loc = plot_opts['legend_loc']

    xlabel = plot_opts['xlabel']
    ylabel = plot_opts['ylabel']
    ylabel = "%s"%(ylabel.replace('%','\%'))
    xlim = plot_opts['xlim']
    ylim = plot_opts['ylim']

    title = r"$\bf{%s}$" %(plot_opts['title'])

    matplotlib.rcParams.update({'font.size': font_size})
    matplotlib.rcParams.update({'axes.titlesize': font_size_axis})
    matplotlib.rcParams.update({'axes.titleweight': 'black'})
    matplotlib.rcParams.update({'axes.labelsize': font_size_axis})

    fig, ax = plt.subplots()

    index_sort = scores.argsort(descending=False)

    plotted_lines = []
    legend_text = []

    for id, id_sort in enumerate(index_sort):
        line = ax.plot(x.tolist(), y[id_sort, :].tolist(),
                       linewidth=line_width,
                       color=plot_draw_styles[index_sort.numel() - id - 1]['color'],
                       linestyle=plot_draw_styles[index_sort.numel() - id - 1]['line_style'])

        plotted_lines.append(line[0])

        tracker = trackers[id_sort]
        disp_name = get_tracker_display_name(tracker)

        legend_text.append('{} [{:.1f}]'.format(disp_name, scores[id_sort]))

    try:
        # add bold to our method
        for i in range(1,2):
            legend_text[-i] = r'\textbf{%s}'%(legend_text[-i])

        ax.legend(plotted_lines[::-1], legend_text[::-1], loc=legend_loc, fancybox=False, edgecolor='black',
                  fontsize=font_size_legend, framealpha=1.0)
    except:
        pass

    ax.set(xlabel=xlabel,
           ylabel=ylabel,
           xlim=xlim, ylim=ylim,
           title=title)

    ax.grid(True, linestyle='-.')
    fig.tight_layout()

    tikzplotlib.save('{}/{}_plot.tex'.format(result_plot_path, plot_type))
    fig.savefig('{}/{}_plot.pdf'.format(result_plot_path, plot_type), dpi=300, format='pdf', transparent=True)
    plt.draw()


def check_and_load_precomputed_results(trackers, dataset, report_name, force_evaluation=False, **kwargs):
    # Load data
    settings = env_settings()

    # Load pre-computed results
    result_plot_path = os.path.join(settings.result_plot_path, report_name)
    eval_data_path = os.path.join(result_plot_path, 'eval_data.pkl')

    if os.path.isfile(eval_data_path) and not force_evaluation:
        with open(eval_data_path, 'rb') as fh:
            eval_data = pickle.load(fh)
    else:
        # print('Pre-computed evaluation data not found. Computing results!')
        eval_data = extract_results(trackers, dataset, report_name, **kwargs)

    if not check_eval_data_is_valid(eval_data, trackers, dataset):
        # print('Pre-computed evaluation data invalid. Re-computing results!')
        eval_data = extract_results(trackers, dataset, report_name, **kwargs)
        # pass
    else:
        # Update display names
        tracker_names = [{'name': t.name, 'param': t.parameter_name, 'run_id': t.run_id, 'disp_name': t.display_name}
                         for t in trackers]
        eval_data['trackers'] = tracker_names
    with open(eval_data_path, 'wb') as fh:
        pickle.dump(eval_data, fh)
    return eval_data


def get_auc_curve(ave_success_rate_plot_overlap, valid_sequence):
    ave_success_rate_plot_overlap = ave_success_rate_plot_overlap[valid_sequence, :, :]
    auc_curve = ave_success_rate_plot_overlap.mean(0) * 100.0
    auc = auc_curve.mean(-1)

    return auc_curve, auc


def get_prec_curve(ave_success_rate_plot_center, valid_sequence):
    ave_success_rate_plot_center = ave_success_rate_plot_center[valid_sequence, :, :]
    prec_curve = ave_success_rate_plot_center.mean(0) * 100.0
    prec_score = prec_curve[:, 20]

    return prec_curve, prec_score


def plot_results(trackers, dataset, report_name, merge_results=False,
                 plot_types=('success'), force_evaluation=False, **kwargs):
    """
    Plot results for the given trackers

    args:
        trackers - List of trackers to evaluate
        dataset - List of sequences to evaluate
        report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
        merge_results - If True, multiple random runs for a non-deterministic trackers are averaged
        plot_types - List of scores to display. Can contain 'success',
                    'prec' (precision), and 'norm_prec' (normalized precision)
    """
    # Load data
    settings = env_settings()

    plot_draw_styles = get_plot_draw_styles()

    # Load pre-computed results
    result_plot_path = os.path.join(settings.result_plot_path, report_name)
    eval_data = check_and_load_precomputed_results(trackers, dataset, report_name, force_evaluation, **kwargs)

    # Merge results from multiple runs
    if merge_results:
        eval_data = merge_multiple_runs(eval_data)

    tracker_names = eval_data['trackers']

    valid_sequence = torch.tensor(eval_data['valid_sequence'], dtype=torch.bool)

    print('\nPlotting results over {} / {} sequences'.format(valid_sequence.long().sum().item(), valid_sequence.shape[0]))

    print('\nGenerating plots for: {}'.format(report_name))

    # ********************************  Success Plot **************************************
    if 'success' in plot_types:
        ave_success_rate_plot_overlap = torch.tensor(eval_data['ave_success_rate_plot_overlap'])

        # Index out valid sequences
        auc_curve, auc = get_auc_curve(ave_success_rate_plot_overlap, valid_sequence)
        threshold_set_overlap = torch.tensor(eval_data['threshold_set_overlap'])

        success_plot_opts = {'plot_type': 'success', 'legend_loc': 'lower left', 'xlabel': 'Overlap threshold',
                             'ylabel': 'Overlap Precision [%]', 'xlim': (0, 1.0), 'ylim': (0, 88), 'title': 'Success'}
        plot_draw_save(auc_curve, threshold_set_overlap, auc, tracker_names, plot_draw_styles, result_plot_path, success_plot_opts)

    # ********************************  Precision Plot **************************************
    if 'prec' in plot_types:
        ave_success_rate_plot_center = torch.tensor(eval_data['ave_success_rate_plot_center'])

        # Index out valid sequences
        prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center, valid_sequence)
        threshold_set_center = torch.tensor(eval_data['threshold_set_center'])

        precision_plot_opts = {'plot_type': 'precision', 'legend_loc': 'lower right',
                               'xlabel': 'Location error threshold [pixels]', 'ylabel': 'Distance Precision [%]',
                               'xlim': (0, 50), 'ylim': (0, 100), 'title': 'Precision plot'}
        plot_draw_save(prec_curve, threshold_set_center, prec_score, tracker_names, plot_draw_styles, result_plot_path,
                       precision_plot_opts)

    # ********************************  Norm Precision Plot **************************************
    if 'norm_prec' in plot_types:
        ave_success_rate_plot_center_norm = torch.tensor(eval_data['ave_success_rate_plot_center_norm'])

        # Index out valid sequences
        prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center_norm, valid_sequence)
        threshold_set_center_norm = torch.tensor(eval_data['threshold_set_center_norm'])

        norm_precision_plot_opts = {'plot_type': 'norm_precision', 'legend_loc': 'lower right',
                                    'xlabel': 'Location error threshold', 'ylabel': 'Distance Precision [%]',
                                    'xlim': (0, 0.5), 'ylim': (0, 85), 'title': 'Normalized Precision'}
        plot_draw_save(prec_curve, threshold_set_center_norm, prec_score, tracker_names, plot_draw_styles, result_plot_path,
                       norm_precision_plot_opts)

    plt.show()


def generate_formatted_report(row_labels, scores, table_name=''):
    name_width = max([len(d) for d in row_labels] + [len(table_name)]) + 5
    min_score_width = 10

    report_text = '\n{label: <{width}} |'.format(label=table_name, width=name_width)

    score_widths = [max(min_score_width, len(k) + 3) for k in scores.keys()]

    for s, s_w in zip(scores.keys(), score_widths):
        report_text = '{prev} {s: <{width}} |'.format(prev=report_text, s=s, width=s_w)

    report_text = '{prev}\n'.format(prev=report_text)

    for trk_id, d_name in enumerate(row_labels):
        # display name
        report_text = '{prev}{tracker: <{width}} |'.format(prev=report_text, tracker=d_name,
                                                           width=name_width)
        for (score_type, score_value), s_w in zip(scores.items(), score_widths):
            report_text = '{prev} {score: <{width}} |'.format(prev=report_text,
                                                              score='{:0.2f}'.format(score_value[trk_id].item()),
                                                              width=s_w)
        report_text = '{prev}\n'.format(prev=report_text)

    return report_text


def print_results(trackers, dataset, report_name, merge_results=False,
                  plot_types=('success'), **kwargs):
    """ Print the results for the given trackers in a formatted table
    args:
        trackers - List of trackers to evaluate
        dataset - List of sequences to evaluate
        report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
        merge_results - If True, multiple random runs for a non-deterministic trackers are averaged
        plot_types - List of scores to display. Can contain 'success' (prints AUC, OP50, and OP75 scores),
                    'prec' (prints precision score), and 'norm_prec' (prints normalized precision score)
    """
    # Load pre-computed results
    eval_data = check_and_load_precomputed_results(trackers, dataset, report_name, **kwargs)

    # Merge results from multiple runs
    if merge_results:
        eval_data = merge_multiple_runs(eval_data)

    tracker_names = eval_data['trackers']
    valid_sequence = torch.tensor(eval_data['valid_sequence'], dtype=torch.bool)

    print('\nReporting results over {} / {} sequences'.format(valid_sequence.long().sum().item(), valid_sequence.shape[0]))

    scores = {}

    # ********************************  Success Plot **************************************
    if 'success' in plot_types:
        threshold_set_overlap = torch.tensor(eval_data['threshold_set_overlap'])
        ave_success_rate_plot_overlap = torch.tensor(eval_data['ave_success_rate_plot_overlap'])

        # Index out valid sequences
        auc_curve, auc = get_auc_curve(ave_success_rate_plot_overlap, valid_sequence)
        scores['AUC'] = auc
        scores['OP50'] = auc_curve[:, threshold_set_overlap == 0.50]
        scores['OP75'] = auc_curve[:, threshold_set_overlap == 0.75]

    # ********************************  Precision Plot **************************************
    if 'prec' in plot_types:
        ave_success_rate_plot_center = torch.tensor(eval_data['ave_success_rate_plot_center'])

        # Index out valid sequences
        prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center, valid_sequence)
        scores['Precision'] = prec_score

    # ********************************  Norm Precision Plot *********************************
    if 'norm_prec' in plot_types:
        ave_success_rate_plot_center_norm = torch.tensor(eval_data['ave_success_rate_plot_center_norm'])

        # Index out valid sequences
        norm_prec_curve, norm_prec_score = get_prec_curve(ave_success_rate_plot_center_norm, valid_sequence)
        scores['Norm Precision'] = norm_prec_score

    # Print
    tracker_disp_names = [get_tracker_display_name(trk) for trk in tracker_names]
    report_text = generate_formatted_report(tracker_disp_names, scores, table_name=report_name)
    print(report_text)


def plot_got_success(trackers, report_name):
    """ Plot success plot for GOT-10k dataset using the json reports.
    Save the json reports from http://got-10k.aitestunion.com/leaderboard in the directory set to
    env_settings.got_reports_path

    The tracker name in the experiment file should be set to the name of the report file for that tracker,
    e.g. DiMP50_report_2019_09_02_15_44_25 if the report is name DiMP50_report_2019_09_02_15_44_25.json

    args:
        trackers - List of trackers to evaluate
        report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
    """
    # Load data
    settings = env_settings()
    plot_draw_styles = get_plot_draw_styles()

    result_plot_path = os.path.join(settings.result_plot_path, report_name)

    auc_curve = torch.zeros((len(trackers), 101))
    scores = torch.zeros(len(trackers))

    # Load results
    tracker_names = []
    for trk_id, trk in enumerate(trackers):
        json_path = '{}/{}.json'.format(settings.got_reports_path, trk.name)

        if os.path.isfile(json_path):
            with open(json_path, 'r') as f:
                eval_data = json.load(f)
        else:
            raise Exception('Report not found {}'.format(json_path))

        if len(eval_data.keys()) > 1:
            raise Exception

        # First field is the tracker name. Index it out
        eval_data = eval_data[list(eval_data.keys())[0]]
        if 'succ_curve' in eval_data.keys():
            curve = eval_data['succ_curve']
            ao = eval_data['ao']
        elif 'overall' in eval_data.keys() and 'succ_curve' in eval_data['overall'].keys():
            curve = eval_data['overall']['succ_curve']
            ao = eval_data['overall']['ao']
        else:
            raise Exception('Invalid JSON file {}'.format(json_path))

        auc_curve[trk_id, :] = torch.tensor(curve) * 100.0
        scores[trk_id] = ao * 100.0

        tracker_names.append({'name': trk.name, 'param': trk.parameter_name, 'run_id': trk.run_id,
                              'disp_name': trk.display_name})

    threshold_set_overlap = torch.arange(0.0, 1.01, 0.01, dtype=torch.float64)

    success_plot_opts = {'plot_type': 'success', 'legend_loc': 'lower left', 'xlabel': 'Overlap threshold',
                         'ylabel': 'Overlap Precision [%]', 'xlim': (0, 1.0), 'ylim': (0, 100), 'title': 'Success plot'}
    plot_draw_save(auc_curve, threshold_set_overlap, scores, tracker_names, plot_draw_styles, result_plot_path,
                   success_plot_opts)
    plt.show()


def print_per_sequence_results(trackers, dataset, report_name, merge_results=False,
                               filter_criteria=None, **kwargs):
    """ Print per-sequence results for the given trackers. Additionally, the sequences to list can be filtered using
    the filter criteria.

    args:
        trackers - List of trackers to evaluate
        dataset - List of sequences to evaluate
        report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
        merge_results - If True, multiple random runs for a non-deterministic trackers are averaged
        filter_criteria - Filter sequence results which are reported. Following modes are supported
                        None: No filtering. Display results for all sequences in dataset
                        'ao_min': Only display sequences for which the minimum average overlap (AO) score over the
                                  trackers is less than a threshold filter_criteria['threshold']. This mode can
                                  be used to select sequences where at least one tracker performs poorly.
                        'ao_max': Only display sequences for which the maximum average overlap (AO) score over the
                                  trackers is less than a threshold filter_criteria['threshold']. This mode can
                                  be used to select sequences all tracker performs poorly.
                        'delta_ao': Only display sequences for which the performance of different trackers vary by at
                                    least filter_criteria['threshold'] in average overlap (AO) score. This mode can
                                    be used to select sequences where the behaviour of the trackers greatly differ
                                    between each other.
    """
    # Load pre-computed results
    eval_data = check_and_load_precomputed_results(trackers, dataset, report_name, **kwargs)

    # Merge results from multiple runs
    if merge_results:
        eval_data = merge_multiple_runs(eval_data)

    tracker_names = eval_data['trackers']
    valid_sequence = torch.tensor(eval_data['valid_sequence'], dtype=torch.bool)
    sequence_names = eval_data['sequences']
    avg_overlap_all = torch.tensor(eval_data['avg_overlap_all']) * 100.0

    # Filter sequences
    if filter_criteria is not None:
        if filter_criteria['mode'] == 'ao_min':
            min_ao = avg_overlap_all.min(dim=1)[0]
            valid_sequence = valid_sequence & (min_ao < filter_criteria['threshold'])
        elif filter_criteria['mode'] == 'ao_max':
            max_ao = avg_overlap_all.max(dim=1)[0]
            valid_sequence = valid_sequence & (max_ao < filter_criteria['threshold'])
        elif filter_criteria['mode'] == 'delta_ao':
            min_ao = avg_overlap_all.min(dim=1)[0]
            max_ao = avg_overlap_all.max(dim=1)[0]
            valid_sequence = valid_sequence & ((max_ao - min_ao) > filter_criteria['threshold'])
        else:
            raise Exception

    avg_overlap_all = avg_overlap_all[valid_sequence, :]
    sequence_names = [s + ' (ID={})'.format(i) for i, (s, v) in enumerate(zip(sequence_names, valid_sequence.tolist())) if v]

    tracker_disp_names = [get_tracker_display_name(trk) for trk in tracker_names]

    scores_per_tracker = {k: avg_overlap_all[:, i] for i, k in enumerate(tracker_disp_names)}
    report_text = generate_formatted_report(sequence_names, scores_per_tracker)

    print(report_text)


def print_results_per_video(trackers, dataset, report_name, merge_results=False,
                  plot_types=('success'), per_video=False, **kwargs):
    """ Print the results for the given trackers in a formatted table
    args:
        trackers - List of trackers to evaluate
        dataset - List of sequences to evaluate
        report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
        merge_results - If True, multiple random runs for a non-deterministic trackers are averaged
        plot_types - List of scores to display. Can contain 'success' (prints AUC, OP50, and OP75 scores),
                    'prec' (prints precision score), and 'norm_prec' (prints normalized precision score)
    """
    # Load pre-computed results
    eval_data = check_and_load_precomputed_results(trackers, dataset, report_name, **kwargs)

    # Merge results from multiple runs
    if merge_results:
        eval_data = merge_multiple_runs(eval_data)

    seq_lens = len(eval_data['sequences'])
    eval_datas = [{} for _ in range(seq_lens)]
    if per_video:
        for key, value in eval_data.items():
            if len(value) == seq_lens:
                for i in range(seq_lens):
                    eval_datas[i][key] = [value[i]]
            else:
                for i in range(seq_lens):
                    eval_datas[i][key] = value

    tracker_names = eval_data['trackers']
    valid_sequence = torch.tensor(eval_data['valid_sequence'], dtype=torch.bool)

    print('\nReporting results over {} / {} sequences'.format(valid_sequence.long().sum().item(), valid_sequence.shape[0]))

    scores = {}

    # ********************************  Success Plot **************************************
    if 'success' in plot_types:
        threshold_set_overlap = torch.tensor(eval_data['threshold_set_overlap'])
        ave_success_rate_plot_overlap = torch.tensor(eval_data['ave_success_rate_plot_overlap'])

        # Index out valid sequences
        auc_curve, auc = get_auc_curve(ave_success_rate_plot_overlap, valid_sequence)
        scores['AUC'] = auc
        scores['OP50'] = auc_curve[:, threshold_set_overlap == 0.50]
        scores['OP75'] = auc_curve[:, threshold_set_overlap == 0.75]

    # ********************************  Precision Plot **************************************
    if 'prec' in plot_types:
        ave_success_rate_plot_center = torch.tensor(eval_data['ave_success_rate_plot_center'])

        # Index out valid sequences
        prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center, valid_sequence)
        scores['Precision'] = prec_score

    # ********************************  Norm Precision Plot *********************************
    if 'norm_prec' in plot_types:
        ave_success_rate_plot_center_norm = torch.tensor(eval_data['ave_success_rate_plot_center_norm'])

        # Index out valid sequences
        norm_prec_curve, norm_prec_score = get_prec_curve(ave_success_rate_plot_center_norm, valid_sequence)
        scores['Norm Precision'] = norm_prec_score

    # Print
    tracker_disp_names = [get_tracker_display_name(trk) for trk in tracker_names]
    report_text = generate_formatted_report(tracker_disp_names, scores, table_name=report_name)
    print(report_text)

    if per_video:
        for i in range(seq_lens):
            eval_data = eval_datas[i]

            print('\n{} sequences'.format(eval_data['sequences'][0]))

            scores = {}
            valid_sequence = torch.tensor(eval_data['valid_sequence'], dtype=torch.bool)

            # ********************************  Success Plot **************************************
            if 'success' in plot_types:
                threshold_set_overlap = torch.tensor(eval_data['threshold_set_overlap'])
                ave_success_rate_plot_overlap = torch.tensor(eval_data['ave_success_rate_plot_overlap'])

                # Index out valid sequences
                auc_curve, auc = get_auc_curve(ave_success_rate_plot_overlap, valid_sequence)
                scores['AUC'] = auc
                scores['OP50'] = auc_curve[:, threshold_set_overlap == 0.50]
                scores['OP75'] = auc_curve[:, threshold_set_overlap == 0.75]

            # ********************************  Precision Plot **************************************
            if 'prec' in plot_types:
                ave_success_rate_plot_center = torch.tensor(eval_data['ave_success_rate_plot_center'])

                # Index out valid sequences
                prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center, valid_sequence)
                scores['Precision'] = prec_score

            # ********************************  Norm Precision Plot *********************************
            if 'norm_prec' in plot_types:
                ave_success_rate_plot_center_norm = torch.tensor(eval_data['ave_success_rate_plot_center_norm'])

                # Index out valid sequences
                norm_prec_curve, norm_prec_score = get_prec_curve(ave_success_rate_plot_center_norm, valid_sequence)
                scores['Norm Precision'] = norm_prec_score

            # Print
            tracker_disp_names = [get_tracker_display_name(trk) for trk in tracker_names]
            report_text = generate_formatted_report(tracker_disp_names, scores, table_name=report_name)
            print(report_text)

lib/test/analysis/plot_results.py
TsingWei/LiteTrack
import numpy as np
from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
from lib.test.utils.load_text import load_text
import os, json


class antiUAVDataset(BaseDataset):
    """ GOT-10k dataset.

    Publication:
        GOT-10k: A Large High-Diversity Benchmark for Generic Object Tracking in the Wild
        Lianghua Huang, Xin Zhao, and Kaiqi Huang
        arXiv:1810.11981, 2018
        https://arxiv.org/pdf/1810.11981.pdf

    Download dataset from http://got-10k.aitestunion.com/downloads
    """
    def __init__(self, split):
        super().__init__()
        # Split can be test, val, or ltrval (a validation split consisting of videos from the official train set)

        self.base_path = os.path.join(self.env_settings.uav_path, split)

        self.sequence_list = self._get_sequence_list(split)
        self.split = split

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _read_anno(self, anno_path):
        # gt = pandas.read_csv(anno_path, delimiter=',', header=None, dtype=np.float32, na_filter=False, low_memory=False).values
        res_file = os.path.join(anno_path, 'IR_label.json')
        with open(res_file, 'r') as f:
            label_res = json.load(f)
        gt = label_res['gt_rect']
        new_gt = [[0,0,0,0] if len(f)==0 else f for f in gt]
        return np.array(new_gt) 

    def _construct_sequence(self, sequence_name):
        anno_path = '{}/{}'.format(self.base_path, sequence_name)

        # ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64)
        ground_truth_rect = self._read_anno(anno_path)

        frames_path = '{}/{}'.format(self.base_path, sequence_name)
        frame_list = [frame for frame in os.listdir(frames_path) if frame.endswith(".jpg")]
        frame_list.sort(key=lambda f: int(f[:-4]))
        frames_list = [os.path.join(frames_path, frame) for frame in frame_list]

        return Sequence(sequence_name, frames_list, 'got10k', ground_truth_rect.reshape(-1, 4))

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self, split):
        with open('{}/list.txt'.format(self.base_path)) as f:
            sequence_list = f.read().splitlines()
        # if split == 'ltrval':
        #     with open('{}/got10k_val_split.txt'.format(self.env_settings.dataspec_path)) as f:
        #         seq_ids = f.read().splitlines()
        #     sequence_list = [sequence_list[int(x)] for x in seq_ids]
        return sequence_list


lib/test/evaluation/antiuavdataset.py
TsingWei/LiteTrack
from collections import namedtuple
import importlib
from lib.test.evaluation.data import SequenceList

DatasetInfo = namedtuple('DatasetInfo', ['module', 'class_name', 'kwargs'])

pt = "lib.test.evaluation.%sdataset"  # Useful abbreviations to reduce the clutter

dataset_dict = dict(
    otb=DatasetInfo(module=pt % "otb", class_name="OTBDataset", kwargs=dict()),
    nfs=DatasetInfo(module=pt % "nfs", class_name="NFSDataset", kwargs=dict()),
    uav=DatasetInfo(module=pt % "uav", class_name="UAVDataset", kwargs=dict()),
    tc128=DatasetInfo(module=pt % "tc128", class_name="TC128Dataset", kwargs=dict()),
    tc128ce=DatasetInfo(module=pt % "tc128ce", class_name="TC128CEDataset", kwargs=dict()),
    trackingnet=DatasetInfo(module=pt % "trackingnet", class_name="TrackingNetDataset", kwargs=dict()),
    got10k=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='train')),
    got10k_hard=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='hard')),
    got10k_test=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='test')),
    got10k_val=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='val')),
    got10k_ltrval=DatasetInfo(module=pt % "got10k", class_name="GOT10KDataset", kwargs=dict(split='ltrval')),
    lasot=DatasetInfo(module=pt % "lasot", class_name="LaSOTDataset", kwargs=dict()),
    lasot_lmdb=DatasetInfo(module=pt % "lasot_lmdb", class_name="LaSOTlmdbDataset", kwargs=dict()),

    vot18=DatasetInfo(module=pt % "vot", class_name="VOTDataset", kwargs=dict()),
    vot22=DatasetInfo(module=pt % "vot", class_name="VOTDataset", kwargs=dict(year=22)),
    itb=DatasetInfo(module=pt % "itb", class_name="ITBDataset", kwargs=dict()),
    tnl2k=DatasetInfo(module=pt % "tnl2k", class_name="TNL2kDataset", kwargs=dict()),
    lasot_extension_subset=DatasetInfo(module=pt % "lasotextensionsubset", class_name="LaSOTExtensionSubsetDataset",
                                       kwargs=dict()),
    antiuav=DatasetInfo(module=pt % "antiuav", class_name="antiUAVDataset", kwargs=dict(split='train')),
    antiuav_val=DatasetInfo(module=pt % "antiuav", class_name="antiUAVDataset", kwargs=dict(split='validation')),
)


def load_dataset(name: str):
    """ Import and load a single dataset."""
    name = name.lower()
    dset_info = dataset_dict.get(name)
    if dset_info is None:
        raise ValueError('Unknown dataset \'%s\'' % name)

    m = importlib.import_module(dset_info.module)
    dataset = getattr(m, dset_info.class_name)(**dset_info.kwargs)  # Call the constructor
    return dataset.get_sequence_list()


def get_dataset(*args):
    """ Get a single or set of datasets."""
    dset = SequenceList()
    for name in args:
        dset.extend(load_dataset(name))
    return dset
lib/test/evaluation/datasets.py
TsingWei/LiteTrack
import importlib
import os


class EnvSettings:
    def __init__(self):
        test_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))

        self.results_path = '{}/tracking_results/'.format(test_path)
        self.segmentation_path = '{}/segmentation_results/'.format(test_path)
        self.network_path = '{}/networks/'.format(test_path)
        self.result_plot_path = '{}/result_plots/'.format(test_path)
        self.otb_path = ''
        self.nfs_path = ''
        self.uav_path = ''
        self.tpl_path = ''
        self.vot_path = ''
        self.got10k_path = ''
        self.lasot_path = ''
        self.trackingnet_path = ''
        self.davis_dir = ''
        self.youtubevos_dir = ''
        self.antiuav_path = ''

        self.got_packed_results_path = ''
        self.got_reports_path = ''
        self.tn_packed_results_path = ''


def create_default_local_file():
    comment = {'results_path': 'Where to store tracking results',
               'network_path': 'Where tracking networks are stored.'}

    path = os.path.join(os.path.dirname(__file__), 'local.py')
    with open(path, 'w') as f:
        settings = EnvSettings()

        f.write('from test.evaluation.environment import EnvSettings\n\n')
        f.write('def local_env_settings():\n')
        f.write('    settings = EnvSettings()\n\n')
        f.write('    # Set your local paths here.\n\n')

        for attr in dir(settings):
            comment_str = None
            if attr in comment:
                comment_str = comment[attr]
            attr_val = getattr(settings, attr)
            if not attr.startswith('__') and not callable(attr_val):
                if comment_str is None:
                    f.write('    settings.{} = \'{}\'\n'.format(attr, attr_val))
                else:
                    f.write('    settings.{} = \'{}\'    # {}\n'.format(attr, attr_val, comment_str))
        f.write('\n    return settings\n\n')


class EnvSettings_ITP:
    def __init__(self, workspace_dir, data_dir, save_dir):
        self.prj_dir = workspace_dir
        self.save_dir = save_dir
        self.results_path = os.path.join(save_dir, 'test/tracking_results')
        self.segmentation_path = os.path.join(save_dir, 'test/segmentation_results')
        self.network_path = os.path.join(save_dir, 'test/networks')
        self.result_plot_path = os.path.join(save_dir, 'test/result_plots')
        self.otb_path = os.path.join(data_dir, 'otb')
        self.nfs_path = os.path.join(data_dir, 'nfs')
        self.uav_path = os.path.join(data_dir, 'uav')
        self.tc128_path = os.path.join(data_dir, 'TC128')
        self.tpl_path = ''
        self.vot_path = os.path.join(data_dir, 'VOT2019')
        self.got10k_path = os.path.join(data_dir, 'got10k')
        self.got10k_lmdb_path = os.path.join(data_dir, 'got10k_lmdb')
        self.lasot_path = os.path.join(data_dir, 'lasot')
        self.lasot_lmdb_path = os.path.join(data_dir, 'lasot_lmdb')
        self.trackingnet_path = os.path.join(data_dir, 'trackingnet')
        self.vot18_path = os.path.join(data_dir, 'vot2018')
        self.vot22_path = os.path.join(data_dir, 'vot2022')
        self.itb_path = os.path.join(data_dir, 'itb')
        self.tnl2k_path = os.path.join(data_dir, 'tnl2k')
        self.lasot_extension_subset_path_path = os.path.join(data_dir, 'lasot_extension_subset')
        self.davis_dir = ''
        self.youtubevos_dir = ''

        self.got_packed_results_path = ''
        self.got_reports_path = ''
        self.tn_packed_results_path = ''


def create_default_local_file_ITP_test(workspace_dir, data_dir, save_dir):
    comment = {'results_path': 'Where to store tracking results',
               'network_path': 'Where tracking networks are stored.'}

    path = os.path.join(os.path.dirname(__file__), 'local.py')
    with open(path, 'w') as f:
        settings = EnvSettings_ITP(workspace_dir, data_dir, save_dir)

        f.write('from lib.test.evaluation.environment import EnvSettings\n\n')
        f.write('def local_env_settings():\n')
        f.write('    settings = EnvSettings()\n\n')
        f.write('    # Set your local paths here.\n\n')

        for attr in dir(settings):
            comment_str = None
            if attr in comment:
                comment_str = comment[attr]
            attr_val = getattr(settings, attr)
            if not attr.startswith('__') and not callable(attr_val):
                if comment_str is None:
                    f.write('    settings.{} = \'{}\'\n'.format(attr, attr_val))
                else:
                    f.write('    settings.{} = \'{}\'    # {}\n'.format(attr, attr_val, comment_str))
        f.write('\n    return settings\n\n')


def env_settings():
    env_module_name = 'lib.test.evaluation.local'
    try:
        env_module = importlib.import_module(env_module_name)
        return env_module.local_env_settings()
    except:
        env_file = os.path.join(os.path.dirname(__file__), 'local.py')

        # Create a default file
        create_default_local_file()
        raise RuntimeError('YOU HAVE NOT SETUP YOUR local.py!!!\n Go to "{}" and set all the paths you need. '
                           'Then try to run again.'.format(env_file))
lib/test/evaluation/environment.py
TsingWei/LiteTrack
import numpy as np
from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
from lib.test.utils.load_text import load_text
import os


class GOT10KDataset(BaseDataset):
    """ GOT-10k dataset.

    Publication:
        GOT-10k: A Large High-Diversity Benchmark for Generic Object Tracking in the Wild
        Lianghua Huang, Xin Zhao, and Kaiqi Huang
        arXiv:1810.11981, 2018
        https://arxiv.org/pdf/1810.11981.pdf

    Download dataset from http://got-10k.aitestunion.com/downloads
    """
    def __init__(self, split):
        super().__init__()
        # Split can be test, val, or ltrval (a validation split consisting of videos from the official train set)
        if split == 'test' or split == 'val' or split == 'hard':
            self.base_path = os.path.join(self.env_settings.got10k_path, split)
        else:
            self.base_path = os.path.join(self.env_settings.got10k_path, 'train')

        self.sequence_list = self._get_sequence_list(split)
        self.split = split

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _construct_sequence(self, sequence_name):
        anno_path = '{}/{}/groundtruth.txt'.format(self.base_path, sequence_name)

        ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64)

        frames_path = '{}/{}'.format(self.base_path, sequence_name)
        frame_list = [frame for frame in os.listdir(frames_path) if frame.endswith(".jpg")]
        frame_list.sort(key=lambda f: int(f[:-4]))
        frames_list = [os.path.join(frames_path, frame) for frame in frame_list]

        return Sequence(sequence_name, frames_list, 'got10k', ground_truth_rect.reshape(-1, 4))

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self, split):
        with open('{}/list.txt'.format(self.base_path)) as f:
            sequence_list = f.read().splitlines()

        if split == 'ltrval':
            with open('{}/got10k_val_split.txt'.format(self.env_settings.dataspec_path)) as f:
                seq_ids = f.read().splitlines()

            sequence_list = [sequence_list[int(x)] for x in seq_ids]
        return sequence_list

lib/test/evaluation/got10kdataset.py
TsingWei/LiteTrack
import numpy as np
from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
from lib.test.utils.load_text import load_text
import os


class ITBDataset(BaseDataset):
    """ NUS-PRO dataset
    """

    def __init__(self):
        super().__init__()
        self.base_path = self.env_settings.itb_path
        self.sequence_info_list = self._get_sequence_info_list(self.base_path)

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_info_list])

    def _construct_sequence(self, sequence_info):
        sequence_path = sequence_info['path']
        nz = sequence_info['nz']
        ext = sequence_info['ext']
        start_frame = sequence_info['startFrame']
        end_frame = sequence_info['endFrame']

        init_omit = 0
        if 'initOmit' in sequence_info:
            init_omit = sequence_info['initOmit']

        frames = ['{base_path}/{sequence_path}/{frame:0{nz}}.{ext}'.format(base_path=self.base_path,
                                                                           sequence_path=sequence_path, frame=frame_num,
                                                                           nz=nz, ext=ext) for frame_num in
                  range(start_frame + init_omit, end_frame + 1)]

        anno_path = '{}/{}'.format(self.base_path, sequence_info['anno_path'])

        # NOTE: NUS has some weird annos which panda cannot handle
        ground_truth_rect = load_text(str(anno_path), delimiter=(',', None), dtype=np.float64, backend='numpy')
        return Sequence(sequence_info['name'], frames, 'otb', ground_truth_rect[init_omit:, :],
                        object_class=sequence_info['object_class'])

    def __len__(self):
        return len(self.sequence_info_list)

    def get_fileNames(self, rootdir):
        fs = []
        fs_all = []
        for root, dirs, files in os.walk(rootdir, topdown=True):
            files.sort()
            files.sort(key=len)
            if files is not None:
                for name in files:
                    _, ending = os.path.splitext(name)
                    if ending == ".jpg":
                        _, root_ = os.path.split(root)
                        fs.append(os.path.join(root_, name))
                        fs_all.append(os.path.join(root, name))

        return fs_all, fs

    def _get_sequence_info_list(self, base_path):
        sequence_info_list = []
        for scene in os.listdir(base_path):
            if '.' in scene:
                continue
            videos = os.listdir(os.path.join(base_path, scene))
            for video in videos:
                _, fs = self.get_fileNames(os.path.join(base_path, scene, video))
                video_tmp = {"name": video, "path": scene + '/' + video, "startFrame": 1, "endFrame": len(fs),
                             "nz": len(fs[0].split('/')[-1].split('.')[0]), "ext": "jpg",
                             "anno_path": scene + '/' + video + "/groundtruth.txt",
                             "object_class": "unknown"}
                sequence_info_list.append(video_tmp)

        return sequence_info_list  # sequence_info_list_50 #

lib/test/evaluation/itbdataset.py
TsingWei/LiteTrack
import numpy as np
from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
from lib.test.utils.load_text import load_text


class LaSOTExtensionSubsetDataset(BaseDataset):
    """
    LaSOT test set consisting of 280 videos (see Protocol-II in the LaSOT paper)
    Publication:
        LaSOT: A High-quality Large-scale Single Object Tracking Benchmark
        Heng Fan, Hexin Bai, Liting Lin, Fan Yang, Peng Chu, Ge Deng, Sijia Yu, Harshit, Mingzhen Huang, Juehuan Liu,
        Yong Xu, Chunyuan Liao, Lin Yuan, Haibin Ling
        IJCV, 2020
        https://arxiv.org/pdf/2009.03465.pdf
    Download the dataset from http://vision.cs.stonybrook.edu/~lasot/download.html
    """
    def __init__(self):
        super().__init__()
        self.base_path = self.env_settings.lasot_extension_subset_path
        self.sequence_list = self._get_sequence_list()
        self.clean_list = self.clean_seq_list()

    def clean_seq_list(self):
        clean_lst = []
        for i in range(len(self.sequence_list)):
            cls, _ = self.sequence_list[i].split('-')
            clean_lst.append(cls)
        return  clean_lst

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _construct_sequence(self, sequence_name):
        class_name = sequence_name.split('-')[0]
        anno_path = '{}/{}/{}/groundtruth.txt'.format(self.base_path, class_name, sequence_name)

        ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64)

        occlusion_label_path = '{}/{}/{}/full_occlusion.txt'.format(self.base_path, class_name, sequence_name)

        # NOTE: pandas backed seems super super slow for loading occlusion/oov masks
        full_occlusion = load_text(str(occlusion_label_path), delimiter=',', dtype=np.float64, backend='numpy')

        out_of_view_label_path = '{}/{}/{}/out_of_view.txt'.format(self.base_path, class_name, sequence_name)
        out_of_view = load_text(str(out_of_view_label_path), delimiter=',', dtype=np.float64, backend='numpy')

        target_visible = np.logical_and(full_occlusion == 0, out_of_view == 0)

        frames_path = '{}/{}/{}/img'.format(self.base_path, class_name, sequence_name)

        frames_list = ['{}/{:08d}.jpg'.format(frames_path, frame_number) for frame_number in range(1, ground_truth_rect.shape[0] + 1)]

        target_class = class_name
        return Sequence(sequence_name, frames_list, 'lasot_extension_subset', ground_truth_rect.reshape(-1, 4),
                        object_class=target_class, target_visible=target_visible)

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self):
        sequence_list = ['atv-1',
                         'atv-2',
                         'atv-3',
                         'atv-4',
                         'atv-5',
                         'atv-6',
                         'atv-7',
                         'atv-8',
                         'atv-9',
                         'atv-10',
                         'badminton-1',
                         'badminton-2',
                         'badminton-3',
                         'badminton-4',
                         'badminton-5',
                         'badminton-6',
                         'badminton-7',
                         'badminton-8',
                         'badminton-9',
                         'badminton-10',
                         'cosplay-1',
                         'cosplay-10',
                         'cosplay-2',
                         'cosplay-3',
                         'cosplay-4',
                         'cosplay-5',
                         'cosplay-6',
                         'cosplay-7',
                         'cosplay-8',
                         'cosplay-9',
                         'dancingshoe-1',
                         'dancingshoe-2',
                         'dancingshoe-3',
                         'dancingshoe-4',
                         'dancingshoe-5',
                         'dancingshoe-6',
                         'dancingshoe-7',
                         'dancingshoe-8',
                         'dancingshoe-9',
                         'dancingshoe-10',
                         'footbag-1',
                         'footbag-2',
                         'footbag-3',
                         'footbag-4',
                         'footbag-5',
                         'footbag-6',
                         'footbag-7',
                         'footbag-8',
                         'footbag-9',
                         'footbag-10',
                         'frisbee-1',
                         'frisbee-2',
                         'frisbee-3',
                         'frisbee-4',
                         'frisbee-5',
                         'frisbee-6',
                         'frisbee-7',
                         'frisbee-8',
                         'frisbee-9',
                         'frisbee-10',
                         'jianzi-1',
                         'jianzi-2',
                         'jianzi-3',
                         'jianzi-4',
                         'jianzi-5',
                         'jianzi-6',
                         'jianzi-7',
                         'jianzi-8',
                         'jianzi-9',
                         'jianzi-10',
                         'lantern-1',
                         'lantern-2',
                         'lantern-3',
                         'lantern-4',
                         'lantern-5',
                         'lantern-6',
                         'lantern-7',
                         'lantern-8',
                         'lantern-9',
                         'lantern-10',
                         'misc-1',
                         'misc-2',
                         'misc-3',
                         'misc-4',
                         'misc-5',
                         'misc-6',
                         'misc-7',
                         'misc-8',
                         'misc-9',
                         'misc-10',
                         'opossum-1',
                         'opossum-2',
                         'opossum-3',
                         'opossum-4',
                         'opossum-5',
                         'opossum-6',
                         'opossum-7',
                         'opossum-8',
                         'opossum-9',
                         'opossum-10',
                         'paddle-1',
                         'paddle-2',
                         'paddle-3',
                         'paddle-4',
                         'paddle-5',
                         'paddle-6',
                         'paddle-7',
                         'paddle-8',
                         'paddle-9',
                         'paddle-10',
                         'raccoon-1',
                         'raccoon-2',
                         'raccoon-3',
                         'raccoon-4',
                         'raccoon-5',
                         'raccoon-6',
                         'raccoon-7',
                         'raccoon-8',
                         'raccoon-9',
                         'raccoon-10',
                         'rhino-1',
                         'rhino-2',
                         'rhino-3',
                         'rhino-4',
                         'rhino-5',
                         'rhino-6',
                         'rhino-7',
                         'rhino-8',
                         'rhino-9',
                         'rhino-10',
                         'skatingshoe-1',
                         'skatingshoe-2',
                         'skatingshoe-3',
                         'skatingshoe-4',
                         'skatingshoe-5',
                         'skatingshoe-6',
                         'skatingshoe-7',
                         'skatingshoe-8',
                         'skatingshoe-9',
                         'skatingshoe-10',
                         'wingsuit-1',
                         'wingsuit-2',
                         'wingsuit-3',
                         'wingsuit-4',
                         'wingsuit-5',
                         'wingsuit-6',
                         'wingsuit-7',
                         'wingsuit-8',
                         'wingsuit-9',
                         'wingsuit-10']
        return sequence_list
lib/test/evaluation/lasotextensionsubsetdataset.py
TsingWei/LiteTrack
import numpy as np
import multiprocessing
import os
import sys
from itertools import product
from collections import OrderedDict
from lib.test.evaluation import Sequence, Tracker
from lib.vis.visdom_cus import Visdom
import torch

VISDOM = None

def _save_tracker_output(seq: Sequence, tracker: Tracker, output: dict):
    """Saves the output of the tracker."""

    if not os.path.exists(tracker.results_dir):
        print("create tracking result dir:", tracker.results_dir)
        os.makedirs(tracker.results_dir)
    if seq.dataset in ['trackingnet', 'got10k']:
        if not os.path.exists(os.path.join(tracker.results_dir, seq.dataset)):
            os.makedirs(os.path.join(tracker.results_dir, seq.dataset))
    '''2021.1.5 create new folder for these two datasets'''
    if seq.dataset in ['trackingnet', 'got10k']:
        base_results_path = os.path.join(tracker.results_dir, seq.dataset, seq.name)
    else:
        base_results_path = os.path.join(tracker.results_dir, seq.name)

    def save_bb(file, data):
        tracked_bb = np.array(data).astype(int)
        np.savetxt(file, tracked_bb, delimiter='\t', fmt='%d')

    def save_time(file, data):
        exec_times = np.array(data).astype(float)
        np.savetxt(file, exec_times, delimiter='\t', fmt='%f')

    def save_score(file, data):
        scores = np.array(data).astype(float)
        np.savetxt(file, scores, delimiter='\t', fmt='%.2f')

    def _convert_dict(input_dict):
        data_dict = {}
        for elem in input_dict:
            for k, v in elem.items():
                if k in data_dict.keys():
                    data_dict[k].append(v)
                else:
                    data_dict[k] = [v, ]
        return data_dict

    for key, data in output.items():
        # If data is empty
        if not data:
            continue

        if key == 'target_bbox':
            if isinstance(data[0], (dict, OrderedDict)):
                data_dict = _convert_dict(data)

                for obj_id, d in data_dict.items():
                    bbox_file = '{}_{}.txt'.format(base_results_path, obj_id)
                    save_bb(bbox_file, d)
            else:
                # Single-object mode
                bbox_file = '{}.txt'.format(base_results_path)
                save_bb(bbox_file, data)

        if key == 'all_boxes':
            if isinstance(data[0], (dict, OrderedDict)):
                data_dict = _convert_dict(data)

                for obj_id, d in data_dict.items():
                    bbox_file = '{}_{}_all_boxes.txt'.format(base_results_path, obj_id)
                    save_bb(bbox_file, d)
            else:
                # Single-object mode
                bbox_file = '{}_all_boxes.txt'.format(base_results_path)
                save_bb(bbox_file, data)

        if key == 'all_scores':
            if isinstance(data[0], (dict, OrderedDict)):
                data_dict = _convert_dict(data)

                for obj_id, d in data_dict.items():
                    bbox_file = '{}_{}_all_scores.txt'.format(base_results_path, obj_id)
                    save_score(bbox_file, d)
            else:
                # Single-object mode
                print("saving scores...")
                bbox_file = '{}_all_scores.txt'.format(base_results_path)
                save_score(bbox_file, data)

        elif key == 'time':
            if isinstance(data[0], dict):
                data_dict = _convert_dict(data)

                for obj_id, d in data_dict.items():
                    timings_file = '{}_{}_time.txt'.format(base_results_path, obj_id)
                    save_time(timings_file, d)
            else:
                timings_file = '{}_time.txt'.format(base_results_path)
                save_time(timings_file, data)


def run_sequence(seq: Sequence, tracker: Tracker, debug=False, num_gpu=8, vis=None):
    """Runs a tracker on a sequence."""
    '''2021.1.2 Add multiple gpu support'''
    try:
        worker_name = multiprocessing.current_process().name
        worker_id = int(worker_name[worker_name.find('-') + 1:]) - 1
        gpu_id = worker_id % num_gpu
        torch.cuda.set_device(gpu_id)
    except:
        pass

    def _results_exist():
        if seq.object_ids is None:
            if seq.dataset in ['trackingnet', 'got10k']:
                base_results_path = os.path.join(tracker.results_dir, seq.dataset, seq.name)
                bbox_file = '{}.txt'.format(base_results_path)
            else:
                bbox_file = '{}/{}.txt'.format(tracker.results_dir, seq.name)
            return os.path.isfile(bbox_file)
        else:
            bbox_files = ['{}/{}_{}.txt'.format(tracker.results_dir, seq.name, obj_id) for obj_id in seq.object_ids]
            missing = [not os.path.isfile(f) for f in bbox_files]
            return sum(missing) == 0

    if _results_exist() and not debug:
        print('FPS: {}'.format(-1))
        return

    print('Tracker: {} {} {} ,  Sequence: {}'.format(tracker.name, tracker.parameter_name, tracker.run_id, seq.name))

    if debug:
        output = tracker.run_sequence(seq, debug=debug, vis=vis)
    else:
        try:
            output = tracker.run_sequence(seq, debug=debug)
            
        except Exception as e:
            print(e)
            return
    
    sys.stdout.flush()

    if isinstance(output['time'][0], (dict, OrderedDict)):
        exec_time = sum([sum(times.values()) for times in output['time']])
        num_frames = len(output['time'])
    else:
        exec_time = sum(output['time'])
        num_frames = len(output['time'])

    print('FPS: {}'.format(num_frames / exec_time))

    if not debug:
        _save_tracker_output(seq, tracker, output)
    if tracker is not None:
        del tracker


def run_dataset(dataset, trackers, debug=False, threads=0, num_gpus=8):
    global VISDOM
    """Runs a list of trackers on a dataset.
    args:
        dataset: List of Sequence instances, forming a dataset.
        trackers: List of Tracker params instances.
        debug: Debug level.
        threads: Number of threads to use (default 0).
    """
    multiprocessing.set_start_method('spawn', force=True)

    print('Evaluating {:4d} trackers on {:5d} sequences'.format(len(trackers), len(dataset)))

    multiprocessing.set_start_method('spawn', force=True)

    if threads == 0:
        mode = 'sequential'
    else:
        mode = 'parallel'

    if mode == 'sequential':
        for seq in dataset:
            for tracker_info in trackers:
                tracker_name, tracker_param, dataset_name, ep_id = tracker_info
                tracker = Tracker(tracker_name, tracker_param, dataset_name, ep_id, debug=debug)
                # VISDOM.register_event_handler(tracker.tracker._visdom_ui_handler, 'Tracking')
                if debug:
                    VISDOM = Visdom(1,{},{})
                run_sequence(seq, tracker, debug=debug, vis=VISDOM)
                # del tracker
    elif mode == 'parallel':
        param_list = [(seq, Tracker(*tracker_info), debug, num_gpus) for seq, tracker_info in product(dataset, trackers)]
        with multiprocessing.Pool(processes=threads) as pool:
            pool.starmap(run_sequence, param_list)
    print('Done')

lib/test/evaluation/running.py
TsingWei/LiteTrack
import os

import numpy as np
from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
from lib.test.utils.load_text import load_text, load_str

############
# current 00000492.png of test_015_Sord_video_Q01_done is damaged and replaced by a copy of 00000491.png
############


class TNL2kDataset(BaseDataset):
    """
    TNL2k test set
    """
    def __init__(self):
        super().__init__()
        self.base_path = self.env_settings.tnl2k_path
        self.sequence_list = self._get_sequence_list()

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _construct_sequence(self, sequence_name):
        # class_name = sequence_name.split('-')[0]
        anno_path = '{}/{}/groundtruth.txt'.format(self.base_path, sequence_name)

        ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64)

        text_dsp_path = '{}/{}/language.txt'.format(self.base_path, sequence_name)
        text_dsp = load_str(text_dsp_path)

        frames_path = '{}/{}/imgs'.format(self.base_path, sequence_name)
        frames_list = [f for f in os.listdir(frames_path)]
        frames_list = sorted(frames_list)
        frames_list = ['{}/{}'.format(frames_path, frame_i) for frame_i in frames_list]

        # target_class = class_name
        return Sequence(sequence_name, frames_list, 'tnl2k', ground_truth_rect.reshape(-1, 4), text_dsp=text_dsp)

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self):
        sequence_list = []
        for seq in os.listdir(self.base_path):
            if os.path.isdir(os.path.join(self.base_path, seq)):
                sequence_list.append(seq)

        return sequence_list

lib/test/evaluation/tnl2kdataset.py
TsingWei/LiteTrack
import importlib
import os
from collections import OrderedDict
from lib.test.evaluation.environment import env_settings
import time
import cv2 as cv

# from lib.utils.lmdb_utils import decode_img
from pathlib import Path
import numpy as np


def trackerlist(name: str, parameter_name: str, dataset_name: str, run_ids = None, display_name: str = None,
                result_only=False):
    """Generate list of trackers.
    args:
        name: Name of tracking method.
        parameter_name: Name of parameter file.
        run_ids: A single or list of run_ids.
        display_name: Name to be displayed in the result plots.
    """
    if run_ids is None or isinstance(run_ids, int):
        run_ids = [run_ids]
    return [Tracker(name, parameter_name, dataset_name, run_id, display_name, result_only) for run_id in run_ids]


class Tracker:
    """Wraps the tracker for evaluation and running purposes.
    args:
        name: Name of tracking method.
        parameter_name: Name of parameter file.
        run_id: The run id.
        display_name: Name to be displayed in the result plots.
    """

    def __init__(self, name: str, parameter_name: str, dataset_name: str, run_id: int = None, display_name: str = None,
                 result_only=False, debug=False):
        assert run_id is None or isinstance(run_id, int)

        self.name = name
        self.parameter_name = parameter_name
        self.dataset_name = dataset_name
        self.run_id = run_id
        self.display_name = display_name

        env = env_settings()
        if self.run_id is None:
            self.results_dir = '{}/{}/{}'.format(env.results_path, self.name, self.parameter_name)
        else:
            self.results_dir = '{}/{}/{}_{:03d}'.format(env.results_path, self.name, self.parameter_name, self.run_id)
        if result_only:
            self.results_dir = '{}/{}'.format(env.results_path, self.name)

        tracker_module_abspath = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                                              '..', 'tracker', '%s.py' % self.name))
        if os.path.isfile(tracker_module_abspath):
            tracker_module = importlib.import_module('lib.test.tracker.{}'.format(self.name))
            self.tracker_class = tracker_module.get_tracker_class()
        else:
            self.tracker_class = None
        params = self.get_parameters()
        debug_ = debug
        if debug is None:
            debug_ = getattr(params, 'debug', 0)
        params.debug = debug_
        self.params = params
        # self.create_tracker(params)

    def create_tracker(self, params):
        self.tracker = self.tracker_class(params, self.dataset_name)

    def run_sequence(self, seq, debug=None, vis=None):
        """Run tracker on sequence.
        args:
            seq: Sequence to run the tracker on.
            visualization: Set visualization flag (None means default value specified in the parameters).
            debug: Set debug level (None means default value specified in the parameters).
            multiobj_mode: Which mode to use for multiple objects.
        """
        # Get init information
        self.create_tracker(self.params)
        init_info = seq.init_info()
        output = self._track_sequence( seq, init_info, vis)
        out = output.copy()
        del self.tracker
        return out

    def _track_sequence(self, seq, init_info, vis=None):
        # Define outputs
        # Each field in output is a list containing tracker prediction for each frame.

        # In case of single object tracking mode:
        # target_bbox[i] is the predicted bounding box for frame i
        # time[i] is the processing time for frame i

        # In case of multi object tracking mode:
        # target_bbox[i] is an OrderedDict, where target_bbox[i][obj_id] is the predicted box for target obj_id in
        # frame i
        # time[i] is either the processing time for frame i, or an OrderedDict containing processing times for each
        # object in frame i

        output = {'target_bbox': [],
                  'time': []}
        if self.tracker.params.save_all_boxes:
            output['all_boxes'] = []
            output['all_scores'] = []

        def _store_outputs(tracker_out: dict, defaults=None):
            defaults = {} if defaults is None else defaults
            for key in output.keys():
                val = tracker_out.get(key, defaults.get(key, None))
                if key in tracker_out or val is not None:
                    output[key].append(val)

        # Initialize
        image = self._read_image(seq.frames[0])

        start_time = time.time()
        out = self.tracker.initialize(image, init_info)
        if out is None:
            out = {}

        prev_output = OrderedDict(out)
        init_default = {'target_bbox': init_info.get('init_bbox'),
                        'time': time.time() - start_time}
        if self.tracker.params.save_all_boxes:
            init_default['all_boxes'] = out['all_boxes']
            init_default['all_scores'] = out['all_scores']

        _store_outputs(out, init_default)

        for frame_num, frame_path in enumerate(seq.frames[1:], start=1):
            image = self._read_image(frame_path)

            start_time = time.time()

            info = seq.frame_info(frame_num)
            info['previous_output'] = prev_output

            if len(seq.ground_truth_rect) > 1:
                info['gt_bbox'] = seq.ground_truth_rect[frame_num]
            out = self.tracker.track(image, info, vis)
            prev_output = OrderedDict(out)
            _store_outputs(out, {'time': time.time() - start_time})

        for key in ['target_bbox', 'all_boxes', 'all_scores']:
            if key in output and len(output[key]) <= 1:
                output.pop(key)

        return output

    def run_video(self, videofilepath, optional_box=None, debug=None, visdom_info=None, save_results=False):
        """Run the tracker with the vieofile.
        args:
            debug: Debug level.
        """

        params = self.get_parameters()

        debug_ = debug
        if debug is None:
            debug_ = getattr(params, 'debug', 0)
        params.debug = debug_

        params.tracker_name = self.name
        params.param_name = self.parameter_name
        # self._init_visdom(visdom_info, debug_)

        multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default'))

        if multiobj_mode == 'default':
            self.create_tracker(params)
        else:
            raise ValueError('Unknown multi object mode {}'.format(multiobj_mode))

        assert os.path.isfile(videofilepath), "Invalid param {}".format(videofilepath)
        ", videofilepath must be a valid videofile"

        output_boxes = []

        cap = cv.VideoCapture(videofilepath)
        display_name = 'Display: ' + self.tracker.params.tracker_name
        cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO)
        cv.resizeWindow(display_name, 960, 720)
        success, frame = cap.read()
        cv.imshow(display_name, frame)

        def _build_init_info(box):
            return {'init_bbox': box}

        if success is not True:
            print("Read frame from {} failed.".format(videofilepath))
            exit(-1)
        if optional_box is not None:
            assert isinstance(optional_box, (list, tuple))
            assert len(optional_box) == 4, "valid box's foramt is [x,y,w,h]"
            self.tracker.initialize(frame, _build_init_info(optional_box))
            output_boxes.append(optional_box)
        else:
            while True:
                # cv.waitKey()
                frame_disp = frame.copy()

                cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL,
                           1.5, (0, 0, 0), 1)

                x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False)
                init_state = [x, y, w, h]
                self.tracker.initialize(frame, _build_init_info(init_state))
                output_boxes.append(init_state)
                break

        while True:
            ret, frame = cap.read()

            if frame is None:
                break

            frame_disp = frame.copy()

            # Draw box
            out = self.tracker.track(frame)
            state = [int(s) for s in out['target_bbox']]
            output_boxes.append(state)

            cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]),
                         (0, 255, 0), 5)

            font_color = (0, 0, 0)
            cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1,
                       font_color, 1)
            cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1,
                       font_color, 1)
            cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1,
                       font_color, 1)

            # Display the resulting frame
            cv.imshow(display_name, frame_disp)
            key = cv.waitKey(1)
            if key == ord('q'):
                break
            elif key == ord('r'):
                ret, frame = cap.read()
                frame_disp = frame.copy()

                cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5,
                           (0, 0, 0), 1)

                cv.imshow(display_name, frame_disp)
                x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False)
                init_state = [x, y, w, h]
                self.tracker.initialize(frame, _build_init_info(init_state))
                output_boxes.append(init_state)

        # When everything done, release the capture
        cap.release()
        cv.destroyAllWindows()

        if save_results:
            if not os.path.exists(self.results_dir):
                os.makedirs(self.results_dir)
            video_name = Path(videofilepath).stem
            base_results_path = os.path.join(self.results_dir, 'video_{}'.format(video_name))

            tracked_bb = np.array(output_boxes).astype(int)
            bbox_file = '{}.txt'.format(base_results_path)
            np.savetxt(bbox_file, tracked_bb, delimiter='\t', fmt='%d')


    def get_parameters(self):
        """Get parameters."""
        param_module = importlib.import_module('lib.test.parameter.{}'.format(self.name))
        params = param_module.parameters(self.parameter_name, self.run_id)
        return params

    def _read_image(self, image_file: str):
        if isinstance(image_file, str):
            im = cv.imread(image_file)
            return cv.cvtColor(im, cv.COLOR_BGR2RGB)
        # elif isinstance(image_file, list) and len(image_file) == 2:
        #     return decode_img(image_file[0], image_file[1])
        else:
            raise ValueError("type of image_file should be str or list")




lib/test/evaluation/tracker.py
TsingWei/LiteTrack
from typing import Union, TextIO

import numpy as np
from numba import jit

from lib.test.evaluation.data import SequenceList, BaseDataset, Sequence


class VOTDataset(BaseDataset):
    """
    VOT2018 dataset

    Publication:
        The sixth Visual Object Tracking VOT2018 challenge results.
        Matej Kristan, Ales Leonardis, Jiri Matas, Michael Felsberg, Roman Pfugfelder, Luka Cehovin Zajc, Tomas Vojir,
        Goutam Bhat, Alan Lukezic et al.
        ECCV, 2018
        https://prints.vicos.si/publications/365

    Download the dataset from http://www.votchallenge.net/vot2018/dataset.html
    """
    def __init__(self, year=18):
        super().__init__()
        self.year = year
        if year == 18:
            self.base_path = self.env_settings.vot18_path
        elif year == 20:
            self.base_path = self.env_settings.vot20_path
        elif year == 22:
            self.base_path = self.env_settings.vot22_path
        self.sequence_list = self._get_sequence_list(year)

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _construct_sequence(self, sequence_name):
        sequence_path = sequence_name
        nz = 8
        ext = 'jpg'
        start_frame = 1

        anno_path = '{}/{}/groundtruth.txt'.format(self.base_path, sequence_name)

        if self.year == 18 or self.year == 22:
            try:
                ground_truth_rect = np.loadtxt(str(anno_path), dtype=np.float64)
            except:
                ground_truth_rect = np.loadtxt(str(anno_path), delimiter=',', dtype=np.float64)

            end_frame = ground_truth_rect.shape[0]

            frames = ['{base_path}/{sequence_path}/color/{frame:0{nz}}.{ext}'.format(base_path=self.base_path,
                      sequence_path=sequence_path, frame=frame_num, nz=nz, ext=ext)
                      for frame_num in range(start_frame, end_frame+1)]

            # Convert gt
            if ground_truth_rect.shape[1] > 4:
                gt_x_all = ground_truth_rect[:, [0, 2, 4, 6]]
                gt_y_all = ground_truth_rect[:, [1, 3, 5, 7]]

                x1 = np.amin(gt_x_all, 1).reshape(-1,1)
                y1 = np.amin(gt_y_all, 1).reshape(-1,1)
                x2 = np.amax(gt_x_all, 1).reshape(-1,1)
                y2 = np.amax(gt_y_all, 1).reshape(-1,1)

                ground_truth_rect = np.concatenate((x1, y1, x2-x1, y2-y1), 1)

        elif self.year == 20:
            ground_truth_rect = read_file(str(anno_path))
            ground_truth_rect = np.array(ground_truth_rect, dtype=np.float64)
            end_frame = ground_truth_rect.shape[0]

            frames = ['{base_path}/{sequence_path}/color/{frame:0{nz}}.{ext}'.format(base_path=self.base_path,
                                                                                     sequence_path=sequence_path,
                                                                                     frame=frame_num, nz=nz, ext=ext)
                      for frame_num in range(start_frame, end_frame + 1)]

        else:
            raise NotImplementedError

        return Sequence(sequence_name, frames, 'vot', ground_truth_rect)

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self, year):
        if year == 18:
            sequence_list= ['ants1',
                            'ants3',
                            'bag',
                            'ball1',
                            'ball2',
                            'basketball',
                            'birds1',
                            'blanket',
                            'bmx',
                            'bolt1',
                            'bolt2',
                            'book',
                            'butterfly',
                            'car1',
                            'conduction1',
                            'crabs1',
                            'crossing',
                            'dinosaur',
                            'drone_across',
                            'drone_flip',
                            'drone1',
                            'fernando',
                            'fish1',
                            'fish2',
                            'fish3',
                            'flamingo1',
                            'frisbee',
                            'girl',
                            'glove',
                            'godfather',
                            'graduate',
                            'gymnastics1',
                            'gymnastics2',
                            'gymnastics3',
                            'hand',
                            'handball1',
                            'handball2',
                            'helicopter',
                            'iceskater1',
                            'iceskater2',
                            'leaves',
                            'matrix',
                            'motocross1',
                            'motocross2',
                            'nature',
                            'pedestrian1',
                            'rabbit',
                            'racing',
                            'road',
                            'shaking',
                            'sheep',
                            'singer2',
                            'singer3',
                            'soccer1',
                            'soccer2',
                            'soldier',
                            'tiger',
                            'traffic',
                            'wiper',
                            'zebrafish1']
        elif year == 20:

            sequence_list= ['agility',
                            'ants1',
                            'ball2',
                            'ball3',
                            'basketball',
                            'birds1',
                            'bolt1',
                            'book',
                            'butterfly',
                            'car1',
                            'conduction1',
                            'crabs1',
                            'dinosaur',
                            'dribble',
                            'drone1',
                            'drone_across',
                            'drone_flip',
                            'fernando',
                            'fish1',
                            'fish2',
                            'flamingo1',
                            'frisbee',
                            'girl',
                            'glove',
                            'godfather',
                            'graduate',
                            'gymnastics1',
                            'gymnastics2',
                            'gymnastics3',
                            'hand',
                            'hand02',
                            'hand2',
                            'handball1',
                            'handball2',
                            'helicopter',
                            'iceskater1',
                            'iceskater2',
                            'lamb',
                            'leaves',
                            'marathon',
                            'matrix',
                            'monkey',
                            'motocross1',
                            'nature',
                            'polo',
                            'rabbit',
                            'rabbit2',
                            'road',
                            'rowing',
                            'shaking',
                            'singer2',
                            'singer3',
                            'soccer1',
                            'soccer2',
                            'soldier',
                            'surfing',
                            'tiger',
                            'wheel',
                            'wiper',
                            'zebrafish1']
        elif year == 22:
            sequence_list= ['agility',
                            'animal',
                            'ants1',
                            'bag',
                            'ball2',
                            'ball3',
                            'basketball',
                            'birds1',
                            'birds2',
                            'bolt1',
                            'book',
                            'bubble',
                            'butterfly',
                            'car1',
                            'conduction1',
                            'crabs1',
                            'dinosaur',
                            'diver',
                            'drone1',
                            'drone_across',
                            'fernando',
                            'fish1',
                            'fish2',
                            'flamingo1',
                            'frisbee',
                            'girl',
                            'graduate',
                            'gymnastics1',
                            'gymnastics2',
                            'gymnastics3',
                            'hand',
                            'hand2',
                            'handball1',
                            'handball2',
                            'helicopter',
                            'iceskater1',
                            'iceskater2',
                            'kangaroo',
                            'lamb',
                            'leaves',
                            'marathon',
                            'matrix',
                            'monkey',
                            'motocross1',
                            'nature',
                            'polo',
                            'rabbit',
                            'rabbit2',
                            'rowing',
                            'shaking',
                            'singer2',
                            'singer3',
                            'snake',
                            'soccer1',
                            'soccer2',
                            'soldier',
                            'surfing',
                            'tennis',
                            'tiger',
                            'wheel',
                            'wiper',
                            'zebrafish1']

        else:
            raise NotImplementedError

        return sequence_list


def parse(string):
    """
    parse string to the appropriate region format and return region object
    """
    from vot.region.shapes import Rectangle, Polygon, Mask


    if string[0] == 'm':
        # input is a mask - decode it
        m_, offset_, region = create_mask_from_string(string[1:].split(','))
        # return Mask(m_, offset=offset_)
        return region
    else:
        # input is not a mask - check if special, rectangle or polygon
        raise NotImplementedError
    print('Unknown region format.')
    return None


def read_file(fp: Union[str, TextIO]):
    if isinstance(fp, str):
        with open(fp) as file:
            lines = file.readlines()
    else:
        lines = fp.readlines()

    regions = []
    # iterate over all lines in the file
    for i, line in enumerate(lines):
        regions.append(parse(line.strip()))
    return regions


def create_mask_from_string(mask_encoding):
    """
    mask_encoding: a string in the following format: x0, y0, w, h, RLE
    output: mask, offset
    mask: 2-D binary mask, size defined in the mask encoding
    offset: (x, y) offset of the mask in the image coordinates
    """
    elements = [int(el) for el in mask_encoding]
    tl_x, tl_y, region_w, region_h = elements[:4]
    rle = np.array([el for el in elements[4:]], dtype=np.int32)

    # create mask from RLE within target region
    mask = rle_to_mask(rle, region_w, region_h)
    region = [tl_x, tl_y, region_w, region_h]

    return mask, (tl_x, tl_y), region

@jit(nopython=True)
def rle_to_mask(rle, width, height):
    """
    rle: input rle mask encoding
    each evenly-indexed element represents number of consecutive 0s
    each oddly indexed element represents number of consecutive 1s
    width and height are dimensions of the mask
    output: 2-D binary mask
    """
    # allocate list of zeros
    v = [0] * (width * height)

    # set id of the last different element to the beginning of the vector
    idx_ = 0
    for i in range(len(rle)):
        if i % 2 != 0:
            # write as many 1s as RLE says (zeros are already in the vector)
            for j in range(rle[i]):
                v[idx_+j] = 1
        idx_ += rle[i]
lib/test/evaluation/votdataset.py
TsingWei/LiteTrack
from lib.test.utils import TrackerParams
import os
from lib.test.evaluation.environment import env_settings
from lib.config.litetrack.config import cfg, update_config_from_file


def parameters(yaml_name: str, run_epoch):
    params = TrackerParams()
    prj_dir = env_settings().prj_dir
    save_dir = env_settings().save_dir
    # update default config from yaml file
    yaml_file = os.path.join(prj_dir, 'experiments/litetrack/%s.yaml' % yaml_name)
    update_config_from_file(yaml_file)
    params.cfg = cfg
    # print("test config: ", cfg)

    # template and search region
    params.template_factor = cfg.TEST.TEMPLATE_FACTOR
    params.template_size = cfg.TEST.TEMPLATE_SIZE
    params.search_factor = cfg.TEST.SEARCH_FACTOR
    params.search_size = cfg.TEST.SEARCH_SIZE

    # Network checkpoint path
    params.checkpoint = os.path.join(save_dir, "checkpoints/train/litetrack/%s/LiteTrack_ep%04d.pth.tar" %
                                     (yaml_name, run_epoch))

    # whether to save boxes from all queries
    params.save_all_boxes = False

    return params

lib/test/parameter/litetrack.py
TsingWei/LiteTrack
from lib.test.utils import TrackerParams
import os
from lib.test.evaluation.environment import env_settings
from lib.config.ostrack.config import cfg, update_config_from_file


def parameters(yaml_name: str, run_epoch):
    params = TrackerParams()
    prj_dir = env_settings().prj_dir
    save_dir = env_settings().save_dir
    # update default config from yaml file
    yaml_file = os.path.join(prj_dir, 'experiments/ostrack/%s.yaml' % yaml_name)
    update_config_from_file(yaml_file)
    params.cfg = cfg
    # print("test config: ", cfg)

    # template and search region
    params.template_factor = cfg.TEST.TEMPLATE_FACTOR
    params.template_size = cfg.TEST.TEMPLATE_SIZE
    params.search_factor = cfg.TEST.SEARCH_FACTOR
    params.search_size = cfg.TEST.SEARCH_SIZE

    # Network checkpoint path
    params.checkpoint = os.path.join(save_dir, "checkpoints/train/ostrack/%s/OSTrack_ep%04d.pth.tar" %
                                     (yaml_name, run_epoch))

    # whether to save boxes from all queries
    params.save_all_boxes = False

    return params

lib/test/parameter/ostrack.py
TsingWei/LiteTrack
import time

import torch
from _collections import OrderedDict

from lib.train.data.processing_utils import transform_image_to_crop
from lib.vis.visdom_cus import Visdom


class BaseTracker:
    """Base class for all trackers."""

    def __init__(self, params):
        self.params = params
        self.visdom = None
        self.pause_mode = False
        self.step = False
        self.next_seq = False

    def predicts_segmentation_mask(self):
        return False

    def initialize(self, image, info: dict) -> dict:
        """Overload this function in your tracker. This should initialize the model."""
        raise NotImplementedError

    def track(self, image, info: dict = None) -> dict:
        """Overload this function in your tracker. This should track in the frame and update the model."""
        raise NotImplementedError

    def visdom_draw_tracking(self, image, box, segmentation=None):
        if isinstance(box, OrderedDict):
            box = [v for k, v in box.items()]
        else:
            box = (box,)
        if segmentation is None:
            self.visdom.register((image, *box), 'Tracking', 1, 'Tracking')
        else:
            self.visdom.register((image, *box, segmentation), 'Tracking', 1, 'Tracking')

    def transform_bbox_to_crop(self, box_in, resize_factor, device, box_extract=None, crop_type='template'):
        # box_in: list [x1, y1, w, h], not normalized
        # box_extract: same as box_in
        # out bbox: Torch.tensor [1, 1, 4], x1y1wh, normalized
        if crop_type == 'template':
            crop_sz = torch.Tensor([self.params.template_size, self.params.template_size])
        elif crop_type == 'search':
            crop_sz = torch.Tensor([self.params.search_size, self.params.search_size])
        else:
            raise NotImplementedError

        box_in = torch.tensor(box_in)
        if box_extract is None:
            box_extract = box_in
        else:
            box_extract = torch.tensor(box_extract)
        template_bbox = transform_image_to_crop(box_in, box_extract, resize_factor, crop_sz, normalize=True)
        template_bbox = template_bbox.view(1, 1, 4).to(device)

        return template_bbox

    # def _init_visdom(self, visdom_info, debug):
    #     visdom_info = {} if visdom_info is None else visdom_info
    #     if debug > 0 and visdom_info.get('use_visdom', True):
    #         # try:
    #         self.visdom = Visdom(debug, {'handler': self._visdom_ui_handler, 'win_id': 'Tracking'},
    #                                 visdom_info=visdom_info)

    #             # # Show help
    #             # help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \
    #             #             'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \
    #             #             'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \
    #             #             'block list.'
    #             # self.visdom.register(help_text, 'text', 1, 'Help')
    #         # except:
    #         #     time.sleep(0.5)
    #         #     print('!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n'
    #         #           '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!')

    def _visdom_ui_handler(self, data):
        if data['event_type'] == 'KeyPress':
            if data['key'] == ' ':
                self.pause_mode = not self.pause_mode

            elif data['key'] == 'ArrowRight' and self.pause_mode:
                self.step = True

            elif data['key'] == 'n':
                self.next_seq = True

lib/test/tracker/basetracker.py
TsingWei/LiteTrack
import math

from lib.models import build_LiteTrack
from lib.test.tracker.basetracker import BaseTracker
import torch

from lib.test.tracker.vis_utils import gen_visualization
from lib.test.utils.hann import hann2d
from lib.train.data.processing_utils import sample_target
# for debug
import cv2
import os
import numpy as np

from lib.test.tracker.data_utils import Preprocessor
from lib.utils.box_ops import clip_box, box_xyxy_to_cxcywh, box_xyxy_to_xywh, box_xywh_to_xyxy, box_cxcywh_to_xywh



class LiteTrack(BaseTracker):
    def __init__(self, params, dataset_name):
        super(LiteTrack, self).__init__(params)
        network = build_LiteTrack(params.cfg, training=False)
        network.load_state_dict(torch.load(self.params.checkpoint, map_location='cpu')['net'], strict=False)

        self.cfg = params.cfg
        self.network = network.cuda()
        self.network.eval()
        self.preprocessor = Preprocessor()
        self.state = None

        self.feat_sz = self.cfg.TEST.SEARCH_SIZE // self.cfg.MODEL.BACKBONE.STRIDE
        # motion constrain
        self.output_window = hann2d(torch.tensor([self.feat_sz, self.feat_sz]).long(), centered=True).cuda()

        # for debug
        self.debug = params.debug
        self.use_visdom = 0
        self.frame_id = 0
        if self.debug:
            if not self.use_visdom:
                self.save_dir = "debug"
                if not os.path.exists(self.save_dir):
                    os.makedirs(self.save_dir)
            else:
                # self.add_hook()
                # self._init_visdom(None, 1)
                pass
        # for save boxes from all queries
        self.save_all_boxes = params.save_all_boxes
        self.z_dict1 = {}

    def initialize(self, image, info: dict):
        # forward the template once
        z_patch_arr, resize_factor, z_amask_arr = sample_target(image, info['init_bbox'], self.params.template_factor,
                                                    output_sz=self.params.template_size)
        self.z_patch_arr = z_patch_arr
        template = self.preprocessor.process(z_patch_arr, z_amask_arr)

        self.box_mask_z = None
        
        template_bbox = self.transform_bbox_to_crop(info['init_bbox'], resize_factor,
                                                    template.tensors.device).squeeze(1)
        template_bbox = box_xywh_to_xyxy(template_bbox).float()

        with torch.no_grad():
            self.z_dict1 = template
            self.z_feat = self.network.forward_z(template.tensors, template_bb=template_bbox)
            # self.z_dict1 = self.network.forward_template(template.tensors, self.template_target_box_mask_z)

        # save states
        self.state = info['init_bbox']
        
        # self.template_bbbox = torch.ones_like(template_bbox).float().cuda()
        self.frame_id = 0
        self.last_update_frame = self.frame_id
        if self.save_all_boxes:
            '''save all predicted boxes'''
            all_boxes_save = info['init_bbox'] * self.cfg.MODEL.NUM_OBJECT_QUERIES
            return {"all_boxes": all_boxes_save}

    def track(self, image, info: dict = None, vis=None):
        if self.cfg.MODEL.HEAD.TYPE == 'DECODER':
            return self.track_decoder(image, info=info, vis=vis)
        elif self.cfg.MODEL.HEAD.TYPE == 'GFL':
            return self.track_GFL(image, info=info, vis=vis)
        elif self.cfg.MODEL.HEAD.TYPE == 'CENTER':
            return self.track_center(image, info=info, vis=vis)

    def track_center(self, image, info: dict = None, vis=None):
        H, W, _ = image.shape
        self.frame_id += 1
        x_patch_arr, resize_factor, x_amask_arr = sample_target(image, self.state, self.params.search_factor,
                                                                output_sz=self.params.search_size)  # (x1, y1, w, h)
        search = self.preprocessor.process(x_patch_arr, x_amask_arr)

        with torch.no_grad():
            x_dict = search
            # merge the template and the search
            # run the transformer
            # out_dict = self.network.forward(
            #     template=self.z_dict1.tensors, search=x_dict.tensors, template_bb=self.template_bbox)
            out_dict = self.network(
                template_feats=self.z_feat, search=x_dict.tensors)    

        # add hann windows
        pred_score_map = out_dict['score_map']
        response = self.output_window * pred_score_map
        pred_boxes = self.network.box_head.cal_bbox(response, int(self.feat_sz),  out_dict['size_map'], out_dict['offset_map'])
        pred_boxes = pred_boxes.view(-1, 4)
        # Baseline: Take the mean of all pred boxes as the final result
        pred_box = (pred_boxes.mean(
            dim=0) * self.params.search_size / resize_factor).tolist()  # (cx, cy, w, h) [0,1]
        # get the final box result
        self.state = clip_box(self.map_box_back(pred_box, resize_factor), H, W, margin=10)

        # for debug
        if self.debug:
            if not self.use_visdom:
                x1, y1, w, h = self.state
                image_BGR = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                cv2.rectangle(image_BGR, (int(x1),int(y1)), (int(x1+w),int(y1+h)), color=(0,0,255), thickness=2)
                save_path = os.path.join(self.save_dir, "%04d.jpg" % self.frame_id)
                cv2.imwrite(save_path, image_BGR)
            else:
                vis.register((image, info['gt_bbox'].tolist(), self.state), 'Tracking', 1, 'Tracking')

                vis.register(torch.from_numpy(x_patch_arr).permute(2, 0, 1), 'image', 1, 'search_region')
                vis.register(torch.from_numpy(self.z_patch_arr).permute(2, 0, 1), 'image', 1, 'template')
                vis.register(pred_score_map.view(self.feat_sz, self.feat_sz), 'heatmap', 1, 'score_map')
                vis.register((pred_score_map * self.output_window).view(self.feat_sz, self.feat_sz), 'heatmap', 1, 'score_map_hann')

        if self.save_all_boxes:
            '''save all predictions'''
            all_boxes = self.map_box_back_batch(pred_boxes * self.params.search_size / resize_factor, resize_factor)
            all_boxes_save = all_boxes.view(-1).tolist()  # (4N, )
            return {"target_bbox": self.state,
                    "all_boxes": all_boxes_save}
        else:
            return {"target_bbox": self.state}

    def map_box_back(self, pred_box: list, resize_factor: float):
        cx_prev, cy_prev = self.state[0] + 0.5 * self.state[2], self.state[1] + 0.5 * self.state[3]
        cx, cy, w, h = pred_box
        half_side = 0.5 * self.params.search_size / resize_factor
        cx_real = cx + (cx_prev - half_side)
        cy_real = cy + (cy_prev - half_side)
        return [cx_real - 0.5 * w, cy_real - 0.5 * h, w, h]

    def map_box_back_batch(self, pred_box: torch.Tensor, resize_factor: float):
        cx_prev, cy_prev = self.state[0] + 0.5 * self.state[2], self.state[1] + 0.5 * self.state[3]
        cx, cy, w, h = pred_box.unbind(-1) # (N,4) --> (N,)
        half_side = 0.5 * self.params.search_size / resize_factor
        cx_real = cx + (cx_prev - half_side)
        cy_real = cy + (cy_prev - half_side)
        return torch.stack([cx_real - 0.5 * w, cy_real - 0.5 * h, w, h], dim=-1)

    def add_hook(self):
        conv_features, enc_attn_weights, dec_attn_weights = [], [], []

        for i in range(12):
            self.network.backbone.blocks[i].attn.register_forward_hook(
                # lambda self, input, output: enc_attn_weights.append(output[1])
                lambda self, input, output: enc_attn_weights.append(output[1])
            )

        self.enc_attn_weights = enc_attn_weights


def get_tracker_class():
    return LiteTrack

lib/test/tracker/litetrack.py
TsingWei/LiteTrack
import numpy as np


############## used for visulize eliminated tokens #################
def get_keep_indices(decisions):
    keep_indices = []
    for i in range(3):
        if i == 0:
            keep_indices.append(decisions[i])
        else:
            keep_indices.append(keep_indices[-1][decisions[i]])
    return keep_indices


def gen_masked_tokens(tokens, indices, alpha=0.2):
    # indices = [i for i in range(196) if i not in indices]
    indices = indices[0].astype(int)
    tokens = tokens.copy()
    tokens[indices] = alpha * tokens[indices] + (1 - alpha) * 255
    return tokens


def recover_image(tokens, H, W, Hp, Wp, patch_size):
    # image: (C, 196, 16, 16)
    image = tokens.reshape(Hp, Wp, patch_size, patch_size, 3).swapaxes(1, 2).reshape(H, W, 3)
    return image


def pad_img(img):
    height, width, channels = img.shape
    im_bg = np.ones((height, width + 8, channels)) * 255
    im_bg[0:height, 0:width, :] = img
    return im_bg


def gen_visualization(image, mask_indices, patch_size=16):
    # image [224, 224, 3]
    # mask_indices, list of masked token indices

    # mask mask_indices need to cat
    # mask_indices = mask_indices[::-1]
    num_stages = len(mask_indices)
    for i in range(1, num_stages):
        mask_indices[i] = np.concatenate([mask_indices[i-1], mask_indices[i]], axis=1)

    # keep_indices = get_keep_indices(decisions)
    image = np.asarray(image)
    H, W, C = image.shape
    Hp, Wp = H // patch_size, W // patch_size
    image_tokens = image.reshape(Hp, patch_size, Wp, patch_size, 3).swapaxes(1, 2).reshape(Hp * Wp, patch_size, patch_size, 3)

    stages = [
        recover_image(gen_masked_tokens(image_tokens, mask_indices[i]), H, W, Hp, Wp, patch_size)
        for i in range(num_stages)
    ]
    imgs = [image] + stages
    imgs = [pad_img(img) for img in imgs]
    viz = np.concatenate(imgs, axis=1)
    return viz

lib/test/tracker/vis_utils.py
TsingWei/LiteTrack
import torch
import math
import torch.nn.functional as F


def hann1d(sz: int, centered = True) -> torch.Tensor:
    """1D cosine window."""
    if centered:
        return 0.5 * (1 - torch.cos((2 * math.pi / (sz + 1)) * torch.arange(1, sz + 1).float()))
    w = 0.5 * (1 + torch.cos((2 * math.pi / (sz + 2)) * torch.arange(0, sz//2 + 1).float()))
    return torch.cat([w, w[1:sz-sz//2].flip((0,))])


def hann2d(sz: torch.Tensor, centered = True) -> torch.Tensor:
    """2D cosine window."""
    return hann1d(sz[0].item(), centered).reshape(1, 1, -1, 1) * hann1d(sz[1].item(), centered).reshape(1, 1, 1, -1)


def hann2d_bias(sz: torch.Tensor, ctr_point: torch.Tensor, centered = True) -> torch.Tensor:
    """2D cosine window."""
    distance = torch.stack([ctr_point, sz-ctr_point], dim=0)
    max_distance, _ = distance.max(dim=0)

    hann1d_x = hann1d(max_distance[0].item() * 2, centered)
    hann1d_x = hann1d_x[max_distance[0] - distance[0, 0]: max_distance[0] + distance[1, 0]]
    hann1d_y = hann1d(max_distance[1].item() * 2, centered)
    hann1d_y = hann1d_y[max_distance[1] - distance[0, 1]: max_distance[1] + distance[1, 1]]

    return hann1d_y.reshape(1, 1, -1, 1) * hann1d_x.reshape(1, 1, 1, -1)



def hann2d_clipped(sz: torch.Tensor, effective_sz: torch.Tensor, centered = True) -> torch.Tensor:
    """1D clipped cosine window."""

    # Ensure that the difference is even
    effective_sz += (effective_sz - sz) % 2
    effective_window = hann1d(effective_sz[0].item(), True).reshape(1, 1, -1, 1) * hann1d(effective_sz[1].item(), True).reshape(1, 1, 1, -1)

    pad = (sz - effective_sz) // 2

    window = F.pad(effective_window, (pad[1].item(), pad[1].item(), pad[0].item(), pad[0].item()), 'replicate')

    if centered:
        return window
    else:
        mid = (sz / 2).int()
        window_shift_lr = torch.cat((window[:, :, :, mid[1]:], window[:, :, :, :mid[1]]), 3)
        return torch.cat((window_shift_lr[:, :, mid[0]:, :], window_shift_lr[:, :, :mid[0], :]), 2)


def gauss_fourier(sz: int, sigma: float, half: bool = False) -> torch.Tensor:
    if half:
        k = torch.arange(0, int(sz/2+1))
    else:
        k = torch.arange(-int((sz-1)/2), int(sz/2+1))
    return (math.sqrt(2*math.pi) * sigma / sz) * torch.exp(-2 * (math.pi * sigma * k.float() / sz)**2)


def gauss_spatial(sz, sigma, center=0, end_pad=0):
    k = torch.arange(-(sz-1)/2, (sz+1)/2+end_pad)
    return torch.exp(-1.0/(2*sigma**2) * (k - center)**2)


def label_function(sz: torch.Tensor, sigma: torch.Tensor):
    return gauss_fourier(sz[0].item(), sigma[0].item()).reshape(1, 1, -1, 1) * gauss_fourier(sz[1].item(), sigma[1].item(), True).reshape(1, 1, 1, -1)

def label_function_spatial(sz: torch.Tensor, sigma: torch.Tensor, center: torch.Tensor = torch.zeros(2), end_pad: torch.Tensor = torch.zeros(2)):
    """The origin is in the middle of the image."""
    return gauss_spatial(sz[0].item(), sigma[0].item(), center[0], end_pad[0].item()).reshape(1, 1, -1, 1) * \
           gauss_spatial(sz[1].item(), sigma[1].item(), center[1], end_pad[1].item()).reshape(1, 1, 1, -1)


def cubic_spline_fourier(f, a):
    """The continuous Fourier transform of a cubic spline kernel."""

    bf = (6*(1 - torch.cos(2 * math.pi * f)) + 3*a*(1 - torch.cos(4 * math.pi * f))
           - (6 + 8*a)*math.pi*f*torch.sin(2 * math.pi * f) - 2*a*math.pi*f*torch.sin(4 * math.pi * f)) \
         / (4 * math.pi**4 * f**4)

    bf[f == 0] = 1

    return bf

def max2d(a: torch.Tensor) -> (torch.Tensor, torch.Tensor):
    """Computes maximum and argmax in the last two dimensions."""

    max_val_row, argmax_row = torch.max(a, dim=-2)
    max_val, argmax_col = torch.max(max_val_row, dim=-1)
    argmax_row = argmax_row.view(argmax_col.numel(),-1)[torch.arange(argmax_col.numel()), argmax_col.view(-1)]
    argmax_row = argmax_row.reshape(argmax_col.shape)
    argmax = torch.cat((argmax_row.unsqueeze(-1), argmax_col.unsqueeze(-1)), -1)
    return max_val, argmax

lib/test/utils/hann.py
TsingWei/LiteTrack
import numpy as np
import pandas as pd


def load_text_numpy(path, delimiter, dtype):
    if isinstance(delimiter, (tuple, list)):
        for d in delimiter:
            try:
                ground_truth_rect = np.loadtxt(path, delimiter=d, dtype=dtype)
                return ground_truth_rect
            except:
                pass

        raise Exception('Could not read file {}'.format(path))
    else:
        ground_truth_rect = np.loadtxt(path, delimiter=delimiter, dtype=dtype)
        return ground_truth_rect


def load_text_pandas(path, delimiter, dtype):
    if isinstance(delimiter, (tuple, list)):
        for d in delimiter:
            try:
                ground_truth_rect = pd.read_csv(path, delimiter=d, header=None, dtype=dtype, na_filter=False,
                                                low_memory=False).values
                return ground_truth_rect
            except Exception as e:
                pass

        raise Exception('Could not read file {}'.format(path))
    else:
        ground_truth_rect = pd.read_csv(path, delimiter=delimiter, header=None, dtype=dtype, na_filter=False,
                                        low_memory=False).values
        return ground_truth_rect


def load_text(path, delimiter=' ', dtype=np.float32, backend='numpy'):
    if backend == 'numpy':
        return load_text_numpy(path, delimiter, dtype)
    elif backend == 'pandas':
        return load_text_pandas(path, delimiter, dtype)


def load_str(path):
    with open(path, "r") as f:
        text_str = f.readline().strip().lower()
    return text_str

lib/test/utils/load_text.py
TsingWei/LiteTrack
import numpy as np
import os
import shutil
import argparse
import _init_paths
from lib.test.evaluation.environment import env_settings


def transform_got10k(src_dir, dest_dir):
    
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)
    items = os.listdir(src_dir)
    for item in items:
        if "all" in item:
            continue
        src_path = os.path.join(src_dir, item)
        if "time" not in item:
            seq_name = item.replace(".txt", '')
            seq_dir = os.path.join(dest_dir, seq_name)
            if not os.path.exists(seq_dir):
                os.makedirs(seq_dir)
            new_item = item.replace(".txt", '_001.txt')
            dest_path = os.path.join(seq_dir, new_item)
            bbox_arr = np.loadtxt(src_path, dtype=np.int32, delimiter='\t')
            np.savetxt(dest_path, bbox_arr, fmt='%d', delimiter=',')
        else:
            seq_name = item.replace("_time.txt", '')
            seq_dir = os.path.join(dest_dir, seq_name)
            if not os.path.exists(seq_dir):
                os.makedirs(seq_dir)
            dest_path = os.path.join(seq_dir, item)
            os.system("cp %s %s" % (src_path, dest_path))
    # make zip archive
    shutil.make_archive(src_dir, "zip", src_dir)
    shutil.make_archive(dest_dir, "zip", dest_dir)
    # Remove the original files
    # shutil.rmtree(src_dir)
    shutil.rmtree(dest_dir)


if __name__ == "__main__":
    import glob
    parser = argparse.ArgumentParser(description='transform got10k results.')
    parser.add_argument('--tracker_name', type=str, help='Name of tracking method.')
    parser.add_argument('--cfg_name', type=str, help='Name of config file.')

    args = parser.parse_args()
    tracker_name = args.tracker_name
    cfg_name = args.cfg_name
    env = env_settings()
    result_dir = env.results_path
    src_dir = os.path.join(result_dir, "%s/%s/got10k/" % (tracker_name, cfg_name))
    # dest_dir = os.path.join(result_dir, "%s/got10k_submit/%s/" % (tracker_name, cfg_name))
    
    listglob = glob.glob(src_dir)
    listglob.sort()
    src_list = listglob
    dst_list = [ os.path.join(result_dir, "%s/got10k_submit/%s/"
             % (tracker_name, path.split('/')[-3])) for path in listglob]
    for i in range(len(src_list)):
        transform_got10k(src_list[i], dst_list[i])
    
    # transform_got10k(args.tracker_name, args.cfg_name) for


lib/test/utils/transform_got10k.py
TsingWei/LiteTrack
import os
import sys
env_path = os.path.join(os.path.dirname(__file__), '../../..')
if env_path not in sys.path:
    sys.path.append(env_path)
from lib.test.vot.litetrack_class import run_vot_exp
os.environ['CUDA_VISIBLE_DEVICES'] = '0'


run_vot_exp('litetrack', 'B4_cae_center_all_ep300', ep=300, vis=False, out_conf=True, channel_type='rgb')
lib/test/vot/litetrack.py
TsingWei/LiteTrack
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import pdb
import cv2
import torch
# import vot
import sys
import time
import os
from lib.test.evaluation import Tracker
import lib.test.vot.vot as vot
from lib.test.vot.vot21_utils import *
# from lib.train.dataset.depth_utils import get_rgbd_frame


class vipt(object):
    def __init__(self, tracker_name='', para_name='', run_id=56):
        # create tracker
        tracker_info = Tracker(tracker_name, para_name, "vot22", run_id=run_id)
        params = tracker_info.get_parameters()
        params.visualization = False
        params.debug = False
        # self.tracker = tracker_info.create_tracker(params)
        self.tracker = tracker_info.tracker
        print(tracker_name, self.tracker)

    def write(self, str):
        txt_path = ""
        file = open(txt_path, 'a')
        file.write(str)

    def initialize(self, img_rgb, selection):
        # init on the 1st frame
        x, y, w, h = selection
        bbox = [x,y,w,h]
        self.H, self.W, _ = img_rgb.shape
        init_info = {'init_bbox': bbox}
        _ = self.tracker.initialize(img_rgb, init_info)

    def track(self, img_rgb):
        # track
        outputs = self.tracker.track(img_rgb)
        pred_bbox = outputs['target_bbox']
        # max_score = outputs['best_score']  #.max().cpu().numpy()
        return pred_bbox, None


def run_vot_exp(tracker_name, para_name, ep=180, vis=False, out_conf=False, channel_type='color'):

    torch.set_num_threads(1)
    save_root = os.path.join('', para_name)
    if vis and (not os.path.exists(save_root)):
        os.mkdir(save_root)
    tracker = vipt(tracker_name=tracker_name, para_name=para_name, run_id=ep)

    if channel_type=='rgb':
        channel_type=None
    handle = vot.VOT("rectangle", channels=channel_type)

    selection = handle.region()
    imagefile = handle.frame()
    if not imagefile:
        sys.exit(0)
    if vis:
        '''for vis'''
        seq_name = imagefile.split('/')[-3]
        save_v_dir = os.path.join(save_root,seq_name)
        if not os.path.exists(save_v_dir):
            os.mkdir(save_v_dir)
        cur_time = int(time.time() % 10000)
        save_dir = os.path.join(save_v_dir, str(cur_time))
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

    # read rgbd data
    if isinstance(imagefile, list) and len(imagefile)==2:
        image = get_rgbd_frame(imagefile[0], imagefile[1], dtype='rgbcolormap', depth_clip=True)
    else:
        print(imagefile)
        image = cv2.cvtColor(cv2.imread(imagefile), cv2.COLOR_BGR2RGB) # Right

    tracker.initialize(image, selection)

    while True:
        imagefile = handle.frame()
        if not imagefile:
            break

        # read rgbd data
        if isinstance(imagefile, list) and len(imagefile) == 2:
            image = get_rgbd_frame(imagefile[0], imagefile[1], dtype='rgbcolormap', depth_clip=True)
        else:
            image = cv2.cvtColor(cv2.imread(imagefile), cv2.COLOR_BGR2RGB)  # Right

        b1, max_score = tracker.track(image)

        if out_conf:
            handle.report(vot.Rectangle(*b1), max_score)
        else:
            handle.report(vot.Rectangle(*b1))
        if vis:
            '''Visualization'''
            # original image
            image_ori = image[:,:,::-1].copy() # RGB --> BGR
            image_name = imagefile.split('/')[-1]
            save_path = os.path.join(save_dir, image_name)
            image_b = image_ori.copy()
            cv2.rectangle(image_b, (int(b1[0]), int(b1[1])),
                          (int(b1[0] + b1[2]), int(b1[1] + b1[3])), (0, 0, 255), 2)
            image_b_name = image_name.replace('.jpg','_bbox.jpg')
            save_path = os.path.join(save_dir, image_b_name)
            cv2.imwrite(save_path, image_b)
lib/test/vot/litetrack_class.py
TsingWei/LiteTrack
import imp
from .base_actor import BaseActor
from .ostrack import OSTrackActor
from .dino4track_actor import Dino4TrackActor
from .dab4track_actor import DAB4TrackActor
from .stark_s import STARKSActor
from .bitrack import BiTrackActor
from .detrack import DETRackActor

lib/train/actors/__init__.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor, inverse_sigmoid
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy, box_iou, box_iou_pairwise, xywh_to_cxcywh
import torch
from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap, generate_cls_map, generate_distribution_heatmap, grid_center_2d, grid_center_flattened, bbox2distance, get_2d_gaussian_map
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate
import torch.nn.functional as F


class BiTrackActor(BaseActor):
    """ Actor for training BiTrack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg
        self.search_feat_size = self.cfg.DATA.SEARCH.SIZE//self.cfg.MODEL.BACKBONE.STRIDE
        self.template_feat_size = self.cfg.DATA.TEMPLATE.SIZE//self.cfg.MODEL.BACKBONE.STRIDE

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # mock
        template_bb = data['template_anno'][0]
        search_bb = data['search_anno'][0]
        # data['rand_template_bb'] = jitter_box_in_box(template_bb,5)
        # data['rand_search_bb'] = jitter_box_in_box(search_bb,5)
        
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        assert len(data['search_images']) == 1

        template_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1,
                                                             *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        box_mask_z = None
        ce_keep_rate = None
        template_target_mask = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
                                             data['template_anno'][0], generate_bb_mask_only=False)
        # if self.cfg.MODEL.BACKBONE.CE_LOC:
            # ce_box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
            #                                 data['template_anno'][0])

        #     ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
        #     ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
        #     ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
        #                                         total_epochs=ce_start_epoch + ce_warm_epoch,
        #                                         ITERS_PER_EPOCH=1,
        #                                         base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])
        template_bb = data['template_anno'][0]
        template_bb = box_xywh_to_xyxy(template_bb).clamp(min=0.0,max=1.0)
        search_bb = data['search_anno'][0].clamp(min=0.0,max=1.0)
        search_bb = box_xywh_to_xyxy(search_bb).clamp(min=0.0,max=1.0)
        # rand_template_bb = data['rand_template_bb']
        # rand_search_bb = data['rand_search_bb']
        # template_target_mask = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
        #                                      data['template_anno'][0], generate_bb_mask_only=True)
        # if self.cfg.MODEL.USE_CLS_TOKEN:
        # gt_gaussian_maps_T = generate_heatmap(data['search_anno'], 16)
        # gt_gaussian_maps_T = gt_gaussian_maps_T[-1].unsqueeze(1)
        gt_cls_maps = generate_cls_map(box_xywh_to_xyxy(data['search_anno'][-1]), self.search_feat_size)
        gt_cls_maps_T = generate_cls_map(box_xywh_to_xyxy(data['template_anno'][-1]), self.template_feat_size)

        if len(template_list) == 1:
            template_list = template_list[0]

        out_dict = self.net(template=template_list,
                            search=search_img,
                            return_last_attn=False,
                            template_bb=template_bb,
                            search_box=search_bb,
                            gt_score_map=gt_cls_maps,
                            gt_score_map_T=gt_cls_maps_T,
                            mask_z=template_target_mask,
                            # mock_template_bb=rand_template_bb,
                            # mock_search_bb=rand_search_bb,
                            )

        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        if self.cfg.MODEL.HEAD.TYPE == 'GFL' or self.cfg.MODEL.HEAD.TYPE == 'PyGFL':
            loss, status = self.compute_losses_GFL(pred_dict, gt_dict,)
            loss_T, status_T = self.compute_losses_GFL_T(pred_dict, gt_dict,)
            loss = loss+0.5*loss_T
            status.update(status_T)
        if self.cfg.MODEL.HEAD.TYPE == 'MLP':
            loss, status = self.compute_losses_MLP(pred_dict, gt_dict,)
            # loss_T, status_T = self.compute_losses_GFL_T(pred_dict, gt_dict,)
            # status.update(status_T)
        if return_status:
            return loss, status
        else:
            return loss

    def compute_losses_GFL(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes']
        # pred_center_score_unsig = pred_dict['score_map_c']
        pred_quality_score = pred_dict['score_map_q']
        best_boxes = pred_dict['box_best']
        
        # c_score_unsig = pred_center_score_unsig.permute(0, 2, 3,
        #                               1).reshape(-1, 256)
        q_score = pred_quality_score.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        
        gt_q_map = generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        gt_c_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]

        # if self.cfg.TRAIN.TRAIN_CLS:
        #     gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
        #                         self.search_feat_size,self.search_feat_size)
        #     gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
        #                         self.search_feat_size,self.search_feat_size)
        
        # Duplicate GT box
        bs = gt_q_map.shape[0]
        inside_ind = torch.nonzero(gt_q_map.flatten(1))
        gt_bboxes_multiple = gt_bbox[inside_ind[:,0],:]
        gt_boxes_vec = box_xywh_to_xyxy(gt_bboxes_multiple).clamp(min=0.0,max=1.0)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)

        iou, _ = box_iou(pred_boxes_vec, gt_boxes_vec)
        iou = iou.detach()
        q_score_pos = pred_quality_score.squeeze(1).flatten(1)[inside_ind[:,0], inside_ind[:,0][1]]

        # # Adaptive select the pos sample
        # pos_ind, selected_iou, select_mask = filter_scores(iou, inside_ind, bs=bs) # ([bs_ind, flatten_map_ind])
        # q_score_pos_unsig = q_score_unsig.flatten(1)[pos_ind[:,0], pos_ind[:,0][1]]
        # pred_box_pos = pred_boxes_vec[select_mask]
        # target_box_pos = gt_boxes_vec[select_mask]
        
        # if self.cfg.TRAIN.TRAIN_CLS:
        #     nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
        #     gt_bbox = gt_bbox[nonzero_indices[:,0],:]
        #     best_boxes = best_boxes[nonzero_indices[:,0],:]
        top_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        top_pred_boxes_vec = best_boxes

        giou_loss, _ = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)
        # giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        top_giou, top_iou = self.objective['giou'](top_pred_boxes_vec, top_gt_boxes_vec)  # (BN,4) (BN,4)
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)

        # Map loss
        iou_score = gt_q_map.new_zeros(gt_q_map.flatten(1).shape) 
        iou_score[inside_ind[:,0], inside_ind[:, 1]] = iou

        # Only fetch selected samples.
        pred_ltrb = pred_dict['ltrb_dis'] # already inside the GT box
        pred_ltrb_pos = pred_ltrb.reshape(-1, 17)
        target_ltrb_pos =  bbox2distance(grid_center_2d(gt_q_map, self.search_feat_size),
            gt_boxes_vec * self.search_feat_size, max_dis=self.search_feat_size).reshape(-1)# dont norm to 1

        # The weight for localization, from predicted cls socre.
        target_weight = q_score_pos.detach()

        # Map for classification.
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_dict['search_anno'][0]), self.search_feat_size)
        iou_gaussian = iou_score * gaussian_center_map.flatten(1)
        # target_gaussian_weight = gaussian_center_map.flatten(1)[inside_ind[:,0],inside_ind[:,1]] * target_weight
        # gt_gaussian_maps[-1].unsqueeze(1)

        # Focal loss
        # ctr_loss = self.objective['focal'](c_score_unsig.sigmoid().reshape(bs, -1), gt_c_map.reshape(bs, -1))
        # QFL handle all the points inside the box.
        qf_loss = self.objective['qfl'](inverse_sigmoid(q_score), gt_q_map.reshape(-1), iou_gaussian.reshape(-1))
        # DFL only handle the selected points inside the box
        df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
                
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + qf_loss\
            + df_loss * 0.25\
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = top_iou.detach().mean()
        mean_giou_loss = top_giou.detach().mean()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss.item(),
                #   "Loss/ctr": ctr_loss.item(),
                    "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status

    def compute_losses_GFL_T(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['template_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes_T']
        # pred_center_score_unsig = pred_dict['score_map_c']
        pred_quality_score = pred_dict['score_map_q_T']
        best_boxes = pred_dict['box_best_T']
        
        # c_score_unsig = pred_center_score_unsig.permute(0, 2, 3,
        #                               1).reshape(-1, 256)
        q_score = pred_quality_score.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        
        gt_q_map = generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.template_feat_size)
        gt_c_map = generate_heatmap(gt_dict['template_anno'], self.template_feat_size)[0]

        # if self.cfg.TRAIN.TRAIN_CLS:
        #     gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
        #                         self.search_feat_size,self.search_feat_size)
        #     gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
        #                         self.search_feat_size,self.search_feat_size)
        
        # Duplicate GT box
        bs = gt_q_map.shape[0]
        inside_ind = torch.nonzero(gt_q_map.flatten(1))
        gt_bboxes_multiple = gt_bbox[inside_ind[:,0],:]
        gt_boxes_vec = box_xywh_to_xyxy(gt_bboxes_multiple).clamp(min=0.0,max=1.0)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)

        iou, _ = box_iou(pred_boxes_vec, gt_boxes_vec)
        iou = iou.detach()
        q_score_pos = pred_quality_score.squeeze(1).flatten(1)[inside_ind[:,0], inside_ind[:,0][1]]

        # # Adaptive select the pos sample
        # pos_ind, selected_iou, select_mask = filter_scores(iou, inside_ind, bs=bs) # ([bs_ind, flatten_map_ind])
        # q_score_pos_unsig = q_score_unsig.flatten(1)[pos_ind[:,0], pos_ind[:,0][1]]
        # pred_box_pos = pred_boxes_vec[select_mask]
        # target_box_pos = gt_boxes_vec[select_mask]
        
        # if self.cfg.TRAIN.TRAIN_CLS:
        #     nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
        #     gt_bbox = gt_bbox[nonzero_indices[:,0],:]
        #     best_boxes = best_boxes[nonzero_indices[:,0],:]
        top_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        top_pred_boxes_vec = best_boxes

        giou_loss, _ = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)
        # giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        top_giou, top_iou = self.objective['giou'](top_pred_boxes_vec, top_gt_boxes_vec)  # (BN,4) (BN,4)
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)

        # Map loss
        iou_score = gt_q_map.new_zeros(gt_q_map.flatten(1).shape) 
        iou_score[inside_ind[:,0], inside_ind[:, 1]] = iou

        # Only fetch selected samples.
        pred_ltrb = pred_dict['ltrb_dis_T'] # already inside the GT box
        pred_ltrb_pos = pred_ltrb.reshape(-1, 17)
        target_ltrb_pos =  bbox2distance(grid_center_2d(gt_q_map, self.template_feat_size),
            gt_boxes_vec * self.template_feat_size, max_dis=self.template_feat_size).reshape(-1)# dont norm to 1

        # The weight for localization, from predicted cls socre.
        target_weight = q_score_pos.detach()

        # Map for classification.
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_dict['template_anno'][0]), self.template_feat_size)
        iou_gaussian = iou_score * gaussian_center_map.flatten(1)
        # target_gaussian_weight = gaussian_center_map.flatten(1)[inside_ind[:,0],inside_ind[:,1]] * target_weight
        # gt_gaussian_maps[-1].unsqueeze(1)

        # Focal loss
        # ctr_loss = self.objective['focal'](c_score_unsig.sigmoid().reshape(bs, -1), gt_c_map.reshape(bs, -1))
        # QFL handle all the points inside the box.
        qf_loss = self.objective['qfl'](inverse_sigmoid(q_score), gt_q_map.reshape(-1), iou_gaussian.reshape(-1))
        # DFL only handle the selected points inside the box
        df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
                
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + qf_loss\
            + df_loss * 0.25\
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = top_iou.detach().mean()
        mean_giou_loss = top_giou.detach().mean()
        
        status = {
                    "IoU_T": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    
    def compute_losses_MLP(self, pred_dict, gt_dict):
        K=30
        reg_max=16
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes']
        pred_dist = pred_dict['pred_dist'].reshape(-1, 4*(reg_max+1))
        pred_query_ind = pred_dict['query_ind']
        query_map = torch.zeros([self.bs, self.search_feat_size * self.search_feat_size], device=pred_query_ind.device)
        query_map.scatter_(1, pred_query_ind, 1)
        query_map_2d = query_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        gt_box_map_2d =  generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        query_inbox_map_2d = gt_box_map_2d * query_map_2d

                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        # Grep positive & negative predicts
        pos_box_logits = torch.gather(query_inbox_map_2d.flatten(1), 1, pred_query_ind)
        pos_box_ind = pos_box_logits.reshape(-1).nonzero().reshape(-1)
        neg_box_ind = (pos_box_logits.reshape(-1) == 0).nonzero().reshape(-1)
        pred_box_pos = pred_boxes[pos_box_ind]
        pred_dist_pos = pred_dist[pos_box_ind]
        pred_dist_neg = pred_dist[neg_box_ind]
        
        # Making positve ltrb targets
        inside_ind = torch.nonzero(query_inbox_map_2d.flatten(1))
        gt_box_pos = gt_bbox[inside_ind[:,0],:]
        gt_box_pos_vec = box_xywh_to_xyxy(gt_box_pos).clamp(min=0.0,max=1.0)
        target_ltrb_pos =  bbox2distance(grid_center_2d(query_inbox_map_2d, self.search_feat_size),
            gt_box_pos_vec * self.search_feat_size, max_dis=self.search_feat_size).reshape(-1)# dont norm to 1
        # target_dist_neg = pred_dist_neg.new_zero(pred_dist_neg.shape)
        
        # Box loss
        try:
            giou_loss, _ = self.objective['giou'](pred_box_pos, gt_box_pos_vec)
            giou_loss = giou_loss.mean()
        except:
            giou_loss = torch.as_tensor(0.0).to("cuda")
            _ = torch.as_tensor(0.0).to("cuda")
        l1_loss = self.objective['l1'](pred_box_pos, pred_box_pos)  # (BN,4) (BN,4)

        # Distribution loss
        df_loss = self.objective['dfl'](pred_dist_pos.reshape(-1,17), target_ltrb_pos, weight=None)
        # neg_loss = self.
                
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + df_loss

            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = _.detach().mean()
        # mean_giou_loss = top_giou.detach().mean()
        
        status = {
                    "IoU": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status



def filter_scores(S, Ind, bs, alpha=1.0):
    # Compute the batch indices
    batch_indices = Ind[:, 0]

    # Compute the mean and standard error for each batch
    batch_means = torch.zeros(bs, dtype=torch.float32, device=S.device)
    batch_std = torch.zeros(bs, dtype=torch.float32, device=S.device)
    for i in range(bs):
        mask = (Ind[:, 0] == i)
        batch_scores = S[mask]
        batch_size = batch_scores.size(0)
        if batch_size > 0:
            batch_means[i] = batch_scores.mean()
            batch_std[i] = batch_scores.std() / torch.sqrt(torch.tensor(batch_size))

    # Compute the threshold for each batch
    batch_thresholds = batch_means + alpha * batch_std

    # Compute the mask for the scores that pass the threshold
    # mask = S > batch_thresholds[batch_indices]
    mask = S > 0

    # Filter the indices and scores
    filtered_indices = Ind[mask]
    filtered_scores = S[mask]

    return filtered_indices, filtered_scores, mask

lib/train/actors/bitrack.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy
import torch
from lib.utils.merge import merge_template_search_DAB, merge_template_search
from ...utils.heapmap_utils import generate_heatmap


class DAB4TrackActor(BaseActor):
    """ Actor for training the STARK-S and STARK-ST(Stage1)"""
    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # forward pass
        out_dict = self.forward_pass(data, run_box_head=True, run_cls_head=False)

        # process the groundtruth
        # gt_bboxes = data['search_anno']  # (Ns, batch, 4) (x1,y1,w,h)

        # compute losses
        loss, status = self.compute_losses_center_head(out_dict, data)

        return loss, status

    def forward_pass(self, data, run_box_head, run_cls_head):
        feat_dict_list = []
        # process the templates
        # print("-------------------------HELLO",data.keys())
        # raise Exception("caonima")
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            feat_dict_list.append(self.net(img=NestedTensor(template_img_i, template_att_i), mode='backbone', obj="T"))

        # process the search regions (t-th frame)
        search_img = data['search_images'].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        search_att = data['search_att'].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)
        feat_dict_list.append(self.net(img=NestedTensor(search_img, search_att), mode='backbone', obj="S"))

        # run the transformer and compute losses
        seq_dict = merge_template_search_DAB(feat_dict_list)
        out_dict = self.net(seq_dict=seq_dict, mode="transformer", run_box_head=run_box_head, run_cls_head=run_cls_head)
        # out_dict: (B, N, C), outputs_coord: (1, B, N, C), target_query: (1, B, N, C)
        return out_dict

    def compute_losses_center_head(self, pred_dict, gt_dict, return_status=True):
            # gt gaussian map
            gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
            gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE)
            gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1)

            # Get boxes
            pred_boxes = pred_dict['pred_boxes']
            if torch.isnan(pred_boxes).any():
                raise ValueError("Network outputs is NAN! Stop Training")
            num_queries = pred_boxes.size(1)
            pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
            gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0,
                                                                                                            max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
            # compute giou and iou
            try:
                giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
            except:
                giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
            # compute l1 loss
            l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
            # compute location loss
            if 'score_map' in pred_dict:
                location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps)
            else:
                location_loss = torch.tensor(0.0, device=l1_loss.device)
            # weighted sum
            loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss
            if return_status:
                # status for log
                mean_iou = iou.detach().mean()
                status = {"Loss/total": loss.item(),
                        "Loss/giou": giou_loss.item(),
                        "Loss/l1": l1_loss.item(),
                        "Loss/location": location_loss.item(),
                        "IoU": mean_iou.item(),
                        #   "Coord":pred_boxes_vec[0].detach().cpu()
                        }
                return loss, status
            else:
                return loss

    def compute_losses(self, pred_dict, gt_bbox, return_status=True):
        # Get boxes
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss

    def compute_losses_(self, pred_dict, gt_bbox, return_status=True):
        # Get boxes
        enc_pred_boxes = pred_dict['enc_outputs']['pred_box']
        pred_boxes = pred_dict['pred_boxes']
        all_pred_boxes = torch.cat((pred_boxes,enc_pred_boxes),dim=1)
        if torch.isnan(all_pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        all_num_queries = all_pred_boxes.size(1)
        target_num_queries = pred_boxes.size(1)
        all_pred_boxes_vec = box_cxcywh_to_xyxy(all_pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        target_pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        all_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, all_num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        target_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, target_num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)
        # compute giou and iou
        try:
            _target_giou_loss, target_iou = self.objective['giou'](target_pred_boxes_vec, target_gt_boxes_vec)
            giou_loss, _iou = self.objective['giou'](all_pred_boxes_vec, all_gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, _iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](all_pred_boxes_vec, all_gt_boxes_vec)  # (BN,4) (BN,4)

        # encoder_loss

        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss
        if return_status:
            # status for log
            final_iou = target_iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "IoU": final_iou.item()}
            return loss, status
        else:
            return loss

lib/train/actors/dab4track_actor.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor, inverse_sigmoid
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy, box_iou, box_iou_pairwise, xywh_to_cxcywh, box_xyxy_to_xywh, box_cxcywh_to_xyxy
import torch
from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap, generate_cls_map, generate_distribution_heatmap, grid_center_2d, grid_center_flattened, bbox2distance, get_2d_gaussian_map
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate, generate_mask_cond_search
from ...utils.cdn import prepare_for_cdn
import torch.nn.functional as F
import torchvision

class DETRackActor(BaseActor):
    """ Actor for training DETRack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg
        self.search_feat_size = 20
        self.template_feat_size = 0
        self.count = 0

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        self.count+=1
        if self.count == 12:
            print('a')
            pass
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        assert len(data['search_images']) == 1

        template_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1,
                                                             *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        box_mask_z = None
        ce_keep_rate = None
        # if self.cfg.MODEL.BACKBONE.CE_LOC:
        #     ce_box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
        #                                     data['template_anno'][0])

        #     ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
        #     ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
        #     ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
        #                                         total_epochs=ce_start_epoch + ce_warm_epoch,
        #                                         ITERS_PER_EPOCH=1,
        #                                         base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])
        template_data = data['template_anno'][0]
        template_bb = box_xywh_to_xyxy(template_data).clamp(min=0.0,max=1.0)
        search_data = data['search_anno'][0]
        search_bb = box_xywh_to_xyxy(search_data).clamp(min=0.0,max=1.0)
        # last_ratio = data['last_search_anno_ratio'][0]
        # search_bb = data['search_anno'][0]
        # last_wh = search_bb[:, 2:]*last_ratio
        # last_wh = last_wh.clamp(min=0.0,max=1.0)
        # template_target_mask = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
        #                                      data['template_anno'][0], generate_bb_mask_only=False)
        # search_target_mask = generate_mask_cond_search(self.cfg, template_list[0].shape[0], template_list[0].device,
        #                                      data['search_anno'][0])
        # if self.cfg.MODEL.USE_CLS_TOKEN:


        if len(template_list) == 1:
            template_list = template_list[0]
        
        # gt_cls_maps = generate_cls_map(box_xywh_to_xyxy(data['search_anno'][-1]), self.search_feat_size)
        # if self.cfg.MODEL.HEAD.TYPE == 'REPCENTER':
        #     gt_cls_maps[:,:,:] = 1
        # gt_cls_maps_T = generate_cls_map(box_xywh_to_xyxy(data['template_anno'][-1]), self.template_feat_size)

        # gaussian_center_map = generate_heatmap(data['search_anno'], self.search_feat_size)[0]
        # center_ind_batch, center_ind = torch.where(gaussian_center_map.flatten(1)>=1)
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_img=NestedTensor(template_img_i, template_att_i)
        search_img = data['search_images'].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        search_att = data['search_att'].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)
        
        # feat_dict_list.append(self.net(img=NestedTensor(search_img, search_att), mode='backbone'))
        out_dict = self.net(template=NestedTensor(template_img_i, template_att_i),
                            search=NestedTensor(search_img, search_att),
                            template_bb=template_bb,
                            search_bb=search_bb,
                            mask_z=None,
                            mask_s=None
                            )

        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        if True:
            loss, status = self.compute_losses_DECODER(pred_dict, gt_dict,)
            return loss, status
        else:
            return loss

    def compute_losses_GFL(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes']
        # pred_center_score_unsig = pred_dict['score_map_c']
        pred_quality_score = pred_dict['score_map']
        best_boxes = pred_dict['box_best']
        # qs = pred_dict['qs']
        
        # c_score_unsig = pred_center_score_unsig.permute(0, 2, 3,
        #                               1).reshape(-1, 256)
        q_score = pred_quality_score.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        
        gt_q_map = generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        gt_c_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]

        if self.cfg.TRAIN.TRAIN_CLS:
            gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
        
        # Duplicate GT box
        bs = gt_q_map.shape[0]
        inside_ind = torch.nonzero(gt_q_map.flatten(1))
        gt_bboxes_multiple = gt_bbox[inside_ind[:,0],:]
        gt_boxes_vec = box_xywh_to_xyxy(gt_bboxes_multiple).clamp(min=0.0,max=1.0)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)

        iou, _ = box_iou(pred_boxes_vec, gt_boxes_vec)
        iou = iou.detach()
        q_score_pos = pred_quality_score.squeeze(1).flatten(1)[inside_ind[:,0], inside_ind[:,0][1]]

        # # Adaptive select the pos sample
        # pos_ind, selected_iou, select_mask = filter_scores(iou, inside_ind, bs=bs) # ([bs_ind, flatten_map_ind])
        # q_score_pos_unsig = q_score_unsig.flatten(1)[pos_ind[:,0], pos_ind[:,0][1]]
        # pred_box_pos = pred_boxes_vec[select_mask]
        # target_box_pos = gt_boxes_vec[select_mask]
        
        if self.cfg.TRAIN.TRAIN_CLS:
            nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
            gt_bbox = gt_bbox[nonzero_indices[:,0],:]
            best_boxes = best_boxes[nonzero_indices[:,0],:]
        top_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        top_pred_boxes_vec = best_boxes

        giou_loss, _ = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)
        # giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        top_giou, top_iou = self.objective['giou'](top_pred_boxes_vec, top_gt_boxes_vec)  # (BN,4) (BN,4)
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)

        # Map loss
        iou_score = gt_q_map.new_zeros(gt_q_map.flatten(1).shape) 
        iou_score[inside_ind[:,0], inside_ind[:, 1]] = iou

        # Only fetch selected samples.
        pred_ltrb = pred_dict['ltrb_dis'] # already inside the GT box
        pred_ltrb_pos = pred_ltrb.reshape(-1, 17)
        target_ltrb_pos =  bbox2distance(grid_center_2d(gt_q_map, self.search_feat_size),
            gt_boxes_vec * self.search_feat_size, max_dis=self.search_feat_size).reshape(-1)# dont norm to 1

        # The weight for localization, from predicted cls socre.
        target_weight = q_score_pos.detach()

        # Map for classification.
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_dict['search_anno'][0]), self.search_feat_size)
        iou_gaussian = iou_score * gaussian_center_map.flatten(1)
        # target_gaussian_weight = gaussian_center_map.flatten(1)[inside_ind[:,0],inside_ind[:,1]] * target_weight
        # gt_gaussian_maps[-1].unsqueeze(1)

        # Focal loss
        # ctr_loss = self.objective['focal'](c_score_unsig.sigmoid().reshape(bs, -1), gt_c_map.reshape(bs, -1))
        # QFL handle all the points inside the box.
        # qs_loss = self.objective['qfl'](qs.reshape(-1, 1), gt_q_map.reshape(-1), iou_score.reshape(-1))
        qf_loss = self.objective['qfl'](inverse_sigmoid(q_score), gt_q_map.reshape(-1), iou_gaussian.reshape(-1))
        # DFL only handle the selected points inside the box
        df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
                
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + qf_loss\
            + df_loss * 0.25\
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = top_iou.detach().mean()
        mean_giou_loss = top_giou.detach().mean()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss.item(),
                #   "Loss/ctr": qs_loss.item(),
                    "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    def compute_losses_REPCENTER(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes']
        # pred_center_score_unsig = pred_dict['score_map_c']
        pred_quality_score = pred_dict['score_map']
        best_boxes = pred_dict['box_best']
        # qs = pred_dict['qs']
        
        # c_score_unsig = pred_center_score_unsig.permute(0, 2, 3,
        #                               1).reshape(-1, 256)
        q_score = pred_quality_score.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        
        gt_q_map = generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        # gt_c_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        # gt_q_map[:,:,:] = 1
        if self.cfg.TRAIN.TRAIN_CLS:
            gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            # gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
            #                     self.search_feat_size,self.search_feat_size)
        
        # Duplicate GT box
        bs = gt_q_map.shape[0]
        inside_ind = torch.nonzero(gt_q_map.flatten(1))
        gt_bboxes_multiple = gt_bbox[inside_ind[:,0],:]
        gt_boxes_vec = box_xywh_to_xyxy(gt_bboxes_multiple).clamp(min=0.0,max=1.0)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)

        iou, _ = box_iou(pred_boxes_vec, gt_boxes_vec)
        iou = iou.detach()
        q_score_pos = pred_quality_score.squeeze(1).flatten(1)[inside_ind[:,0], inside_ind[:,0][1]]

        # # Adaptive select the pos sample
        # pos_ind, selected_iou, select_mask = filter_scores(iou, inside_ind, bs=bs) # ([bs_ind, flatten_map_ind])
        # q_score_pos_unsig = q_score_unsig.flatten(1)[pos_ind[:,0], pos_ind[:,0][1]]
        # pred_box_pos = pred_boxes_vec[select_mask]
        # target_box_pos = gt_boxes_vec[select_mask]
        
        if self.cfg.TRAIN.TRAIN_CLS:
            nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
            gt_bbox = gt_bbox[nonzero_indices[:,0],:]
            best_boxes = best_boxes[nonzero_indices[:,0],:]
        top_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        top_pred_boxes_vec = box_cxcywh_to_xyxy(best_boxes)

        giou_loss, _ = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)
        # giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        top_giou, top_iou = self.objective['giou'](top_pred_boxes_vec, top_gt_boxes_vec)  # (BN,4) (BN,4)
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)

        # Map loss
        iou_score = gt_q_map.new_zeros(gt_q_map.flatten(1).shape) 
        iou_score[inside_ind[:,0], inside_ind[:, 1]] = iou

        # Only fetch selected samples.
        # pred_ltrb = pred_dict['ltrb_dis'] # already inside the GT box
        # pred_ltrb_pos = pred_ltrb.reshape(-1, 17)
        target_ltrb_pos =  bbox2distance(grid_center_2d(gt_q_map, self.search_feat_size),
            gt_boxes_vec * self.search_feat_size, max_dis=self.search_feat_size).reshape(-1)# dont norm to 1

        # The weight for localization, from predicted cls socre.
        target_weight = q_score_pos.detach()

        # Map for classification.
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_dict['search_anno'][0]), self.search_feat_size)
        iou_gaussian = iou_score * gaussian_center_map.flatten(1)
        # target_gaussian_weight = gaussian_center_map.flatten(1)[inside_ind[:,0],inside_ind[:,1]] * target_weight
        # gt_gaussian_maps[-1].unsqueeze(1)

        # Focal loss
        # ctr_loss = self.objective['focal'](c_score_unsig.sigmoid().reshape(bs, -1), gt_c_map.reshape(bs, -1))
        # QFL handle all the points inside the box.
        # qs_loss = self.objective['qfl'](qs.reshape(-1, 1), gt_q_map.reshape(-1), iou_score.reshape(-1))
        qf_loss = self.objective['qfl'](inverse_sigmoid(q_score), gt_q_map.reshape(-1), iou_gaussian.reshape(-1))
        # DFL only handle the selected points inside the box
        # df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
                
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + qf_loss\
            # + df_loss * 0.25\
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = top_iou.detach().mean()
        mean_giou_loss = top_giou.detach().mean()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss.item(),
                #   "Loss/ctr": qs_loss.item(),
                    # "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    def compute_losses_REPPOINTS(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gaussian_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        center_ind_batch, center_ind = torch.where(gaussian_center_map.flatten(1)>=1)
        # Get pred
        pred_box_init_vec = pred_dict['init_box']       # xyxy
        pred_box_refine_vec = pred_dict['refine_box']   # xyxy
        pred_cls_score = pred_dict['score_map']
        anchor_box_vec = pred_dict['anchor_box']
        # guess_size = pred_dict['guess_size']
        bs = pred_cls_score.shape[0]

        pred_box_init = pred_box_init_vec.reshape(bs,-1,4)
        pred_box_refine = pred_box_refine_vec.reshape(bs,-1,4)
        anchor_box = anchor_box_vec.reshape(bs,-1,4)
        pred_cls_score = pred_cls_score.permute(0, 2, 3,
                                      1).squeeze(-1) # (n,1)
        pred_cls_score_vec = pred_cls_score.reshape(-1,1)
        num_queries = pred_box_refine.shape[1]
                                           
        if torch.isnan(pred_box_refine).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        # Calc IoU to select the positive samples for loc(refine).
        gt_boxes_vec_all = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4) --> (n,4)
        iou, _ = box_iou(pred_box_init_vec, gt_boxes_vec_all)
        iou = iou.detach()

        pos_inds = torch.where(iou > 0.5)[0]

        pred_box_refine_pos = pred_box_refine_vec[pos_inds]
        gt_box_pos = gt_boxes_vec_all[pos_inds]


        pred_box_init_center = pred_box_init[center_ind_batch, center_ind]
        pred_box_refine_center = pred_box_refine[center_ind_batch, center_ind]
        anchor_box_center = anchor_box[center_ind_batch, center_ind]
        gt_boxes_vec_center = box_xywh_to_xyxy(gt_bbox).view(-1, 4).clamp(min=0.0, max=1.0)


        if self.cfg.TRAIN.TRAIN_CLS:
            gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            # gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
            #                     self.search_feat_size,self.search_feat_size)
        

        
        if self.cfg.TRAIN.TRAIN_CLS:
            nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
            gt_bbox = gt_bbox[nonzero_indices[:,0],:]
            best_boxes = best_boxes[nonzero_indices[:,0],:]


        giou_loss_init_center, iou_init_center = self.objective['giou'](pred_box_init_center, gt_boxes_vec_center)
        giou_loss_refine_center, iou_refine_center = self.objective['giou'](pred_box_refine_center, gt_boxes_vec_center)
        giou_loss_refine_pos, _ = self.objective['giou'](pred_box_refine_pos, gt_box_pos)
        l1_loss_center = self.objective['l1'](pred_box_refine_center, gt_boxes_vec_center)  # (BN,4) (BN,4)
        _, iou_anchor = self.objective['giou'](anchor_box_center, gt_boxes_vec_center)

        a = box_xyxy_to_xywh(anchor_box_center)
        b = box_xyxy_to_xywh(gt_boxes_vec_center)
        # Map loss
        qf_loss = self.objective['qfl'](pred_cls_score_vec, gaussian_center_map.reshape(-1), gaussian_center_map.reshape(-1))
        # DFL only handle the selected points inside the box
        # df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        # giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
        giou_loss_init_center = giou_loss_init_center.mean()
        giou_loss_refine_center = giou_loss_refine_center.mean()
        giou_loss_refine_pos = giou_loss_refine_pos.mean()
                
        # weighted sum
        loss = giou_loss_init_center \
        + giou_loss_refine_center\
        + qf_loss\
        + l1_loss_center * 0.2
            # + df_loss * 0.25\
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = iou_refine_center.detach().mean()
        mean_iou_init = iou_init_center.detach().mean()
        mean_giou_loss = giou_loss_refine_center.detach()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss_center.item(),
                #   "Loss/ctr": qs_loss.item(),
                    # "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                    "IoU_init": mean_iou_init.item(),
                    "IoU_anchor": iou_anchor.mean().item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    
    def compute_losses_DECODER(self, pred_dict, gt_dict):
        # K = 60
        # num_classes = 1
        iou_aware_ce = False
        outputs_without_aux = {
            k: v for k, v in pred_dict.items() if k != "aux_outputs" and k != "enc_outputs"
        }
        # Get init(enc) and final
        pred_query_ind = outputs_without_aux['query_index']
        
        # pred_query_attn = outputs_without_aux['query_attn']
        src_logits = outputs_without_aux['pred_logits']
        _, K, num_classes = src_logits.shape
        # pred_box_init_vec = outputs_without_aux['init_pred_boxes'].view(-1,4)      # xywh
        pred_box_final_vec = outputs_without_aux['pred_boxes'].reshape(-1,4)      # xywh
        # pred_logit_init_vec = outputs_without_aux['init_pred_logits'].reshape(-1, num_classes)
        pred_logit_final_vec = outputs_without_aux['pred_logits'].reshape(-1, num_classes) 
        pred_logit_init_vec = outputs_without_aux['init_pred_logits'].reshape(-1, num_classes)

        # GT
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_box_vec = gt_bbox.unsqueeze(1).repeat(1,K,1).view(-1,4)
        
        # iou_matrix_init, _ = box_iou(box_xywh_to_xyxy(gt_box_vec), box_cxcywh_to_xyxy(pred_box_init_vec))
        # iou_matrix_init_ = iou_matrix_init.reshape(self.bs, -1)
        iou_matrix_final, _ = box_iou(box_xywh_to_xyxy(gt_box_vec), box_cxcywh_to_xyxy(pred_box_final_vec))
        iou_matrix_final_ = iou_matrix_final.reshape(self.bs, -1)
        filtered_pred_init, filtered_gt_init, box_pos_index = filter_predictions_by_iou(gt_bbox, outputs_without_aux['pred_boxes'], iou_threshold=0.7, maxk=1)
        _, _, cls_pos_index_init = filter_predictions_by_iou(gt_bbox, outputs_without_aux['pred_boxes'], iou_threshold=0.5, maxk=1, dynamic=False)
        _, _, cls_pos_index_final = filter_predictions_by_iou(gt_bbox, outputs_without_aux['pred_boxes'], iou_threshold=0.5, maxk=1, dynamic=False)

        # CDN
        # dn_meta = pred_dict['dn_meta']
        # if dn_meta is not None and 'output_known_lbs_bboxes' in dn_meta:
        #     output_known_lbs_bboxes, single_padding, dn_num_group = prep_for_dn(dn_meta)
        #     dn_pred_box = output_known_lbs_bboxes['pred_boxes']
        #     dn_pred_logits = output_known_lbs_bboxes['pred_logits']
        #     dn_gt_box = gt_bbox.unsqueeze(1).repeat(1,dn_num_group*single_padding,1).view(-1,4)
        #     dn_iou, _ = box_iou(box_xywh_to_xyxy(dn_gt_box).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(dn_pred_box.reshape(-1,4)))
        #     # fetch the 1st box of each dn group
        #     dn_pred_box_positive = dn_pred_box.unsqueeze(1).reshape(self.bs, dn_num_group, single_padding, 4)[:, :, 0].reshape(-1,4)
        #     dn_gt_box_pos = gt_bbox.unsqueeze(1).repeat(1,dn_num_group,1).view(-1,4)
        #     dn_gt_labels = torch.linspace(1,0,2)[None].repeat(self.bs, dn_num_group).cuda()
        #     # iou_cdn = 
                                           
        if torch.isnan(pred_box_final_vec).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        query_map = torch.zeros([self.bs, 40*40+20*20+10*10], device=pred_logit_final_vec.device)
        query_map.scatter_(1, pred_query_ind, 1)
        
        # gt_box_map =  generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)

        gt_inbox_labels = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        gt_cls_labels_init = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        gt_cls_labels_final = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        # gt_cls_labels = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.sha67pe[2]+1],
        #                                     dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
        # gt_labels[box_pos_index] = 1
        gt_cls_labels_init[cls_pos_index_init] = 1
        # gt_cls_labels_final[cls_pos_index_final] = 1
        gt_cls_labels_final[cls_pos_index_final,0] = 1
        # pred_box_init_pos =  box_cxcywh_to_xyxy(pred_box_init_vec[box_pos_index]).clamp(min=0.0, max=1.0)
        pred_box_final_pos =  box_cxcywh_to_xyxy(pred_box_final_vec[box_pos_index]).clamp(min=0.0, max=1.0)
        gt_box_pos = box_xywh_to_xyxy(filtered_gt_init).clamp(min=0.0, max=1.0)

        

        # status for log
        scores = outputs_without_aux['pred_logits'].squeeze(dim=2) 
        top_scores, top_indices = torch.max(scores, dim=1)
        # pred_top_ind = torch.gather(pred_query_ind,1,top_indices[:,None])
        # pred_top_map = torch.zeros_like(query_map).scatter_(1, pred_top_ind, 1)
        iou_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, iou_matrix_final.reshape(self.bs, -1))
        # # init_iou_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, iou_matrix_init.reshape(self.bs, -1))
        # pred_cls_map = torch.zeros_like(query_map, dtype=torch.float16).scatter_(1, pred_query_ind, src_logits.squeeze(-1).sigmoid())
        # pred_attn_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, pred_query_attn)
        # topk_iou_map = torch.zeros_like(query_map, dtype=torch.float16).scatter_(1, pred_query_ind, gt_cls_labels.reshape(self.bs, -1))
        # label_cross_map = pred_attn_map * gt_box_map.reshape(self.bs,-1)
        # targets = torch.gather(gt_box_map.flatten(1), 1, pred_query_ind)

        # DEBUG maps
        # pred_top_map = pred_top_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        iou_map_0 = iou_map[:,:1600].view(self.bs, 40, 40)
        iou_map_1 = iou_map[:,1600:2000].view(self.bs, 20, 20)
        iou_map_2 = iou_map[:,2000:2100].view(self.bs, 10, 10) 
        query_map_0 = query_map[:,:1600].view(self.bs, 40, 40)
        query_map_1 = query_map[:,1600:2000].view(self.bs, 20, 20)
        query_map_2 = query_map[:,2000:2100].view(self.bs, 10, 10) 
        # iou_map = iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        # init_iou_map = init_iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        # # topk_iou_map = topk_iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        # pred_cls_map = pred_cls_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        # gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_bbox), self.search_feat_size, alpha=0.2)
        # norm_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        # inbox_topk_map = query_map * gt_box_map
        # box_targets = torch.gather(inbox_topk_map.flatten(1), 1, pred_query_ind)
        # q_in_box_ind = box_targets.reshape(-1).nonzero().reshape(-1)

        # gt_inbox_labels[q_in_box_ind] = 1
        # gt_cls_labels[cls_pos_index] = 1
        # pred_box_init_pos =  box_cxcywh_to_xyxy(pred_box_init_vec[box_pos_index])
        pred_box_final_pos =  box_cxcywh_to_xyxy(pred_box_final_vec[box_pos_index])
        gt_box_pos = box_xywh_to_xyxy(filtered_gt_init).clamp(min=0.0, max=1.0)

        # LOSS
        if iou_aware_ce:
            # ce_loss_init = self.objective['dfocal'](pred_logit_init_vec, iou_matrix_init[:,None].detach())
            # ce_loss_final = self.objective['dfocal'](pred_logit_final_vec, iou_matrix_final[:,None].detach())
            # ce_loss_final = self.objective['qfl'](pred_logit_final_vec, gt_inbox_labels.reshape(-1), iou_matrix_final.detach())
            ce_loss_init = self.objective['qfl'](pred_logit_init_vec, gt_inbox_labels.reshape(-1), iou_matrix_init.detach())
            ce_loss_final = self.objective['dfocal'](pred_logit_final_vec, gt_cls_labels_final)
        else:
            ce_loss_init = self.objective['dfocal'](pred_logit_init_vec, gt_cls_labels_init)
            ce_loss_final = self.objective['dfocal'](pred_logit_final_vec, gt_cls_labels_final)
            # ce_loss_final = self.objective['dfocal'](pred_logit_final_vec.reshape(self.bs, -1).softmax(1).reshape(-1, 1), gt_cls_labels_final)
        
        # ce_loss_final = self.objective['qfl'](pred_logit_final_vec, box_targets.reshape(-1), iou_matrix_final)
        # giou_loss_init, iou_init = self.objective['giou'](pred_box_init_pos, gt_box_pos)
        giou_loss_final, iou_final = self.objective['giou'](pred_box_final_pos, gt_box_pos)
        # l1_loss_init = self.objective['l1'](pred_box_init_pos, gt_box_pos)  # (BN,4) (BN,4)
        l1_loss_final = self.objective['l1'](pred_box_final_pos, gt_box_pos)  # (BN,4) (BN,4)
        # _, iou_anchor = self.objective['giou'](anchor_box_center, gt_boxes_vec_center)

        # CDN LOSS
        if False:
            dn_pred_box_pos =  box_cxcywh_to_xyxy(dn_pred_box_positive)
            dn_gt_box_pos = box_xywh_to_xyxy(dn_gt_box_pos).clamp(min=0.0, max=1.0)
            giou_loss_cdn, iou_cdn_pos = self.objective['giou'](dn_pred_box_pos, dn_gt_box_pos)
            if iou_aware_ce:
                # ce_loss_cdn = self.objective['dfocal'](dn_pred_logits.reshape(-1,1), dn_iou.view(-1,1).detach())
                ce_loss_cdn = self.objective['qfl'](dn_pred_logits.reshape(-1,1), dn_gt_labels.reshape(-1), dn_iou.view(-1,1).detach())
            else:
                ce_loss_cdn = self.objective['dfocal'](dn_pred_logits.reshape(-1,1), dn_gt_labels.view(-1,1))
            ce_loss_cdn = ce_loss_cdn.mean(1).reshape(self.bs,dn_num_group*single_padding).mean(1).sum()
            
            l1_loss_cdn = self.objective['l1'](dn_pred_box_pos, dn_gt_box_pos)  # (BN,4) (BN,4)
            loss_cdn = ce_loss_cdn * 1.0 + giou_loss_cdn.mean()*2.0 + l1_loss_cdn * 5.0
            aux_loss_cdn = 0
            for aux_output in output_known_lbs_bboxes['aux_outputs']:
                _dn_pred_box = aux_output['pred_boxes']
                _boxes_vec = _dn_pred_box.unsqueeze(1).reshape(self.bs, dn_num_group, single_padding, 4)[:, :, 0].reshape(-1,4)
                _logits_vec = aux_output['pred_logits'].reshape(-1, num_classes)
                _boxes_pos = box_cxcywh_to_xyxy(_boxes_vec)
                _giou_loss, _iou = self.objective['giou'](_boxes_pos, dn_gt_box_pos)
                _dn_iou, _ = box_iou(box_xywh_to_xyxy(dn_gt_box).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(_dn_pred_box.reshape(-1,4)))
                if iou_aware_ce:
                    _ce_loss = self.objective['dfocal'](_logits_vec,  _dn_iou.view(-1,1).detach())
                else:
                    _ce_loss = self.objective['dfocal'](_logits_vec, dn_gt_labels.view(-1,1))
                _ce_loss_cdn = _ce_loss.mean(1).reshape(self.bs,dn_num_group*single_padding).mean(1).sum()
                _l1_loss = self.objective['l1'](_boxes_pos, dn_gt_box_pos)  # (BN,4) (BN,4)
                aux_loss_cdn += 1.0 * _ce_loss_cdn + 2.0 * _giou_loss.mean() + 5.0 *_l1_loss
        else:
            loss_cdn = torch.as_tensor(0.0).to("cuda")
            aux_loss_cdn = torch.as_tensor(0.0).to("cuda")

        # AUX LOSS
        aux_loss = 0
        # for aux_output in pred_dict['aux_outputs']:
        #     _logits_vec = aux_output['pred_logits'].view(-1, num_classes)
        #     _boxes_vec = aux_output['pred_boxes'].view(-1,4)
        #     _boxes_pos = box_cxcywh_to_xyxy(_boxes_vec[box_pos_index])
        #     _giou_loss, _ = self.objective['giou'](_boxes_pos, gt_box_pos)
        #     _iou, _ = box_iou(box_xywh_to_xyxy(gt_box_vec).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(_boxes_vec.reshape(-1,4)))
        #     if iou_aware_ce:
        #         # _ce_loss = self.objective['dfocal'](_logits_vec, _iou[:,None].detach())
        #         _ce_loss = self.objective['qfl'](_logits_vec, gt_inbox_labels.reshape(-1), iou_matrix_final.detach())
        #     else:
        #         _ce_loss = self.objective['dfocal'](_logits_vec, gt_cls_labels_final)
        #         # _ce_loss = self.objective['dfocal'](_logits_vec.reshape(self.bs, -1).softmax(1).reshape(-1, 1), gt_cls_labels_final)
        #     _ce_loss = _ce_loss.mean(1).reshape(self.bs,K).mean(1).sum()
        #     _l1_loss = self.objective['l1'](_boxes_pos, gt_box_pos)  # (BN,4) (BN,4)
        #     aux_loss += 1.0 * _ce_loss + 2.0 * _giou_loss.mean() + 5.0 *_l1_loss

        # giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
        # ce_loss_init = ce_loss_init.reshape(self.bs,K).mean(1).mean()
        # ce_loss_final = ce_loss_final.reshape(self.bs,K).sum(1).mean()
        ce_loss_final = ce_loss_final.mean(1).reshape(self.bs,K).mean(1).sum()
        ce_loss_init = ce_loss_init.mean(1).reshape(self.bs,K).mean(1).sum()
        # giou_loss_init = giou_loss_init.mean()
        giou_loss_final = giou_loss_final.mean()

        # init_box_loss = giou_loss_init * 2.0 + l1_loss_init * 5.0
        final_box_loss = giou_loss_final * 2.0 + l1_loss_final * 5.0
                
        # weighted sum
        loss = 0 \
        + ce_loss_init * 1.0\
        + ce_loss_final * 1.0\
        # + init_box_loss * 1.0\
        + final_box_loss * 1.0\
        + aux_loss * 1.0\
        + loss_cdn * 1.0\
        + aux_loss_cdn * 1.0

        
        
        
        # targets = gt_cls_labels.new_zeros((self.bs, 256).shape)
        # for i, m in enumerate(label_cross_map):
        #     if m.sum()<=0:
        #         max_id = gaussian_center_map[i].view(-1).argmax()
        #     else:
        #         max_id = m.view(-1).argmax()
        #     targets[i, max_id] = 1

        top_boxes = torch.gather(outputs_without_aux['pred_boxes'], 1, top_indices[:,None,None].expand(-1, -1, 4)).squeeze(dim=1) 
        pred_iou, _ = box_iou(box_cxcywh_to_xyxy(top_boxes), box_xywh_to_xyxy(gt_bbox).clamp(min=0.0, max=1.0))
        # query_corre_num = (gt_box_map*query_map).sum(1).sum(1)
        # search_img = gt_dict['search_images'][0][7]*torch.asarray([0.229, 0.224, 0.225], device=scores.device)[:,None,None]+\
        #     torch.asarray([0.485, 0.456, 0.406], device=scores.device)[:,None,None]
        # torchvision.utils.save_image(search_img, 'xxxxxx.png')
        # F.interpolate(gt_box_map[None,:], size=(256,256))
        # torchvision.utils.save_image((gt_dict['search_images'][0][2]*torch.asarray([0.229, 0.224, 0.225], device=scores.device)[:,None,None]+\
        #    torch.asarray([0.485, 0.456, 0.406], device=scores.device)[:,None,None])*F.interpolate(gt_box_map[None,:], size=(256,256))[0,2], 'xxxxxx.png')
        mean_iou_pred = pred_iou.detach().mean()
        # mean_iou_init = iou_init.detach().mean()
        mean_iou_last = iou_final.detach().mean()
        cls_loss = ce_loss_final.detach()
        mean_giou_loss = giou_loss_final.detach()
        # avg_box_pos_num = len(box_pos_index)//self.bs
        # avg_cls_pos_num = len(cls_pos_index_final)//self.bs

        status = {  "Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/cls": cls_loss.item(),
                    # "IoU_i": mean_iou_init.item(),
                    "IoU_f": mean_iou_last.item(),
                    "IoU": mean_iou_pred.item(),
                    # "Pb":avg_box_pos_num,
                    # "Pc":avg_cls_pos_num
                    }
        if len(pred_dict['aux_outputs'])>0:
            status['Loss/aux'] = aux_loss
        return loss, status
     
    
def save_masked_img(gt_dict, batch_id, mask_map=None, src='search', f_name='xxx.png', alpha=.2, xyxy=None):
    image_batch = gt_dict[src+'_images'][0,batch_id].cpu()
    std = torch.asarray([0.229, 0.224, 0.225])[:,None,None]
    mean = torch.asarray([0.485, 0.456, 0.406])[:,None,None]
    if src == 'search':
        shape = (320,320)
    elif src == 'template':
        shape = (128,128)
    if mask_map is not None:
        mask = F.interpolate(mask_map[None,:], size=shape)[0, batch_id].cpu() + alpha
    else:
        mask = 1
    img = (image_batch * std + mean)*mask
    if xyxy is not None:
        img_int = (img*256).clamp(0.1, 255.9).to(torch.uint8)
        boxxyxy = xyxy[batch_id] * 320
        # boxxyxy = box_cxcywh_to_xyxy(boxxywh[batch_id]) * 320
        img_int = torchvision.utils.draw_bounding_boxes(img_int, boxxyxy, width=1)
        img = img_int.to(torch.float32)/256.0
    torchvision.utils.save_image(img, f_name)


def filter_scores(S, Ind, bs, alpha=1.0):
    # Compute the batch indices
    batch_indices = Ind[:, 0]

    # Compute the mean and standard error for each batch
    batch_means = torch.zeros(bs, dtype=torch.float32, device=S.device)
    batch_std = torch.zeros(bs, dtype=torch.float32, device=S.device)
    for i in range(bs):
        mask = (Ind[:, 0] == i)
        batch_scores = S[mask]
        batch_size = batch_scores.size(0)
        if batch_size > 0:
            batch_means[i] = batch_scores.mean()
            batch_std[i] = batch_scores.std() / torch.sqrt(torch.tensor(batch_size))

    # Compute the threshold for each batch
    batch_thresholds = batch_means + alpha * batch_std

    # Compute the mask for the scores that pass the threshold
    # mask = S > batch_thresholds[batch_indices]
    mask = S > 0

    # Filter the indices and scores
    filtered_indices = Ind[mask]
    filtered_scores = S[mask]

    return filtered_indices, filtered_scores, mask

   
def batch_box_iou(boxes1, boxes2):
    area1 = boxes1[:, 2] * boxes1[:, 3]
    area2 = boxes2[:, 2] * boxes2[:, 3]

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [bs, N, 2]
    rb = torch.min(boxes1[:, None, 2:4] + boxes1[:, None, :2], boxes2[:, 2:4] + boxes2[:, :2])  # [bs, N, 2]

    wh = (rb - lt).clamp(min=0)  # [bs, N, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [bs, N]

    iou = inter / (area1[:, None] + area2 - inter)
    return iou

def filter_predictions_by_iou(GT_boxes, pred_boxes, iou_threshold=0.5, k=1, maxk=5, dynamic=False):
    bs, n = pred_boxes.shape[:2]

    # Calculate IoU between GT_boxes and pred_boxes
    GT_vec = box_xywh_to_xyxy(GT_boxes)[:, None, :].repeat((1, n, 1)).view(-1, 4).clamp(min=0.0, max=1.0)
    P_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4).clamp(min=0.0, max=1.0)
    iou, _ = box_iou(GT_vec, P_vec)
    iou = iou.view(bs,n)

    # Filter predictions based on IoU threshold
    filtered_preds = []
    filtered_gts = []
    positive_indices = []
    for i in range(bs):
        if dynamic:
            iou_threshold = iou[i].mean()+iou[i].std()
        valid_indices = (iou[i] >= iou_threshold).nonzero(as_tuple=True)[0]
        
        if len(valid_indices) >= maxk:
            topk_indices = iou[i].topk(maxk, largest=True, sorted=True).indices
            valid_indices = topk_indices
        # If no predictions meet the condition, reserve the top-k predictions
        if len(valid_indices) == 0:
            topk_indices = iou[i].topk(k, largest=True, sorted=True).indices
            valid_indices = topk_indices

        filtered_preds.append(pred_boxes[i, valid_indices])
        filtered_gts.append(GT_boxes[i].expand(len(valid_indices), -1))

        # Store the indices of positive samples
        positive_indices.extend((i * n + valid_indices).tolist())

    return torch.cat(filtered_preds), torch.cat(filtered_gts), torch.tensor(positive_indices)

def __filter_predictions_by_iou(gt_boxes, pred_boxes, iou_threshold=0.5, k=1):
    bs, n = pred_boxes.shape[:2]

    # Function to compute IoU between two boxes
    def iou(box1, box2):
        x1, y1, x2, y2 = box1
        X1, Y1, X2, Y2 = box2

        w_intersection = max(0, min(x2, X2) - max(x1, X1))
        h_intersection = max(0, min(y2, Y2) - max(y1, Y1))
        area_intersection = w_intersection * h_intersection

        area_box1 = (x2 - x1) * (y2 - y1)
        area_box2 = (X2 - X1) * (Y2 - Y1)
        area_union = area_box1 + area_box2 - area_intersection

        return area_intersection / area_union

    filtered_pred_boxes = []
    filtered_gt_boxes = []

    for i in range(bs):
        gt_box = gt_boxes[i]
        preds = pred_boxes[i]

        # Compute IoU values between the GT box and predicted boxes
        iou_values = torch.tensor([iou(gt_box, pred) for pred in preds])

        # Filter predictions by the IoU threshold
        filtered_preds = preds[iou_values >= iou_threshold]

        # If no predictions meet the threshold, keep the top-k predictions based on the IoU values
        if len(filtered_preds) == 0:
            topk_indices = iou_values.topk(k, dim=0)[1]
            filtered_preds = preds[topk_indices]

        filtered_pred_boxes.append(filtered_preds)
        filtered_gt_boxes.append(gt_box.unsqueeze(0).repeat(len(filtered_preds),1))

    return filtered_pred_boxes, filtered_gt_boxes

def prep_for_dn(dn_meta):
        output_known_lbs_bboxes = dn_meta['output_known_lbs_bboxes']
        num_dn_groups,pad_size=dn_meta['num_dn_group'],dn_meta['pad_size']
        assert pad_size % num_dn_groups==0
        single_pad=pad_size//num_dn_groups

        return output_known_lbs_bboxes,single_pad,num_dn_groups
lib/train/actors/detrack.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy
import torch
# from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate


class Dino4TrackActor(BaseActor):
    """ Actor for training OSTrack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        assert len(data['search_images']) == 1

        template_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1,
                                                             *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        box_mask_z = None
        ce_keep_rate = None
        # if self.cfg.MODEL.BACKBONE.CE_LOC:
        #     box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
        #                                     data['template_anno'][0])

        #     ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
        #     ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
        #     ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
        #                                         total_epochs=ce_start_epoch + ce_warm_epoch,
        #                                         ITERS_PER_EPOCH=1,
        #                                         base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])

        if len(template_list) == 1:
            template_list = template_list[0]

        out_dict = self.net(template=template_list,
                            search=search_img,
                            # ce_template_mask=box_mask_z,
                            # ce_keep_rate=ce_keep_rate,
                            # return_last_attn=False)
        )

        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        # gt gaussian map
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        # gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE)
        # gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1)

        # Get boxes 
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0,
                                                                                                           max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # compute location loss
        if 'score_map' in pred_dict:
            location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps)
        else:
            location_loss = torch.tensor(0.0, device=l1_loss.device)
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "Loss/location": location_loss.item(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss

lib/train/actors/dino4track_actor.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor, inverse_sigmoid
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy, box_iou, box_iou_pairwise, xywh_to_cxcywh, box_xyxy_to_xywh
import torch
from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap, generate_cls_map, generate_distribution_heatmap, grid_center_2d, grid_center_flattened, bbox2distance, get_2d_gaussian_map
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate, generate_mask_cond_search
from ...utils.cdn import prepare_for_cdn
import torch.nn.functional as F
import torchvision
import torch.distributed as dist

def reduce_mean(tensor):
    if not (dist.is_available() and dist.is_initialized()):
        return tensor
    tensor = tensor.clone()
    dist.all_reduce(tensor.div_(dist.get_world_size()), op=dist.reduce_op.SUM)
    return tensor

class OSTrackActor(BaseActor):
    """ Actor for training OSTrack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg
        self.search_feat_size = self.cfg.DATA.SEARCH.SIZE//self.cfg.MODEL.BACKBONE.STRIDE
        self.template_feat_size = self.cfg.DATA.TEMPLATE.SIZE//self.cfg.MODEL.BACKBONE.STRIDE
        self.count = 0

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        self.count+=1
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        # assert len(data['search_images']) == 1

        template_list = []
        search_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:]).clamp(min=-5.0,max=5.0)  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        for i in range(self.settings.num_search):
            search_img_i = data['search_images'][i].view(-1, *data['search_images'].shape[2:]).clamp(min=-5.0,max=5.0)  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            search_list.append(search_img_i)
        # search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:]).clamp(min=-5.0,max=5.0)  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        box_mask_z = None
        ce_keep_rate = None
        if self.cfg.MODEL.BACKBONE.CE:
            ce_box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
                                            data['template_anno'][0])

            ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
            ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
            ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
                                                total_epochs=ce_start_epoch + ce_warm_epoch,
                                                ITERS_PER_EPOCH=1,
                                                base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])
        template_data = data['template_anno'][0]
        template_bb = box_xywh_to_xyxy(template_data).clamp(min=0.0,max=1.0)
        search_data = data['search_anno'] # [n, bs, 4]
        search_bb = box_xywh_to_xyxy(search_data).clamp(min=0.0,max=1.0)

        # if self.cfg.MODEL.USE_CLS_TOKEN:


        if len(template_list) == 1:
            template_list = template_list[0]
        
        
        gt_cls_maps = generate_cls_map(box_xywh_to_xyxy(data['search_anno'][-1]), self.search_feat_size)
        gt_cls_maps_T = generate_cls_map(box_xywh_to_xyxy(data['template_anno'][-1]), self.template_feat_size).to(torch.int).flatten(1)
        # gt_cls_maps = torch.ones_like(gt_cls_maps)
        # if self.cfg.MODEL.HEAD.TYPE == 'REPCENTER':
        gt_cls_maps[:,:,:] = 1
        gaussian_center_map = generate_heatmap(data['search_anno'], self.search_feat_size)[0]
        if len(search_list) == 1:
            search_list = search_list[0]
            search_bb = search_bb[0]
            out_dict = self.net.forward_train(
                                template=template_list,
                                search=search_list,
                                template_bb=template_bb
                                )
        else:
        # online training
            feat_temperal = None
            out_dict_multiple = []
            for search_img, search_bb_single in zip(search_list, search_bb):
                out_dict = self.net(template=template_list,
                                search=search_img,
                                ce_template_mask=ce_box_mask_z,
                                ce_keep_rate=ce_keep_rate,
                                return_last_attn=False,
                                gt_cls_map=gt_cls_maps,
                                template_bb=template_bb,
                                search_bb=search_bb_single,
                                gt_cls_map_T=None,
                                mask_z=gt_cls_maps_T,
                                mask_s=None,
                                head_only=True,
                                temporal_feat=feat_temperal
                                )
                feat_temperal = out_dict['feat'].detach()
                out_dict_multiple.append(out_dict)
            return out_dict_multiple
        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        if isinstance(pred_dict, list):
            losses=[]
            for i, pred in enumerate(pred_dict):
                gt_dict_single = {k: v[i][None,...] for k, v in gt_dict.items() if k == "search_anno"}
                loss, status = self.compute_losses_DECODER(pred, gt_dict_single,)
                losses.append(loss)
            loss = sum(losses)

        elif self.cfg.MODEL.HEAD.TYPE == 'GFL' or self.cfg.MODEL.HEAD.TYPE == 'PyGFL':
            loss, status = self.compute_losses_GFL(pred_dict, gt_dict,)
        elif self.cfg.MODEL.HEAD.TYPE == 'REPPOINTS':
            loss, status = self.compute_losses_REPPOINTS(pred_dict, gt_dict,)
        elif self.cfg.MODEL.HEAD.TYPE == 'REPCENTER':
            loss, status = self.compute_losses_REPPOINTS_(pred_dict, gt_dict,)
        elif self.cfg.MODEL.HEAD.TYPE == 'DECODER':
            loss, status = self.compute_losses_DECODER(pred_dict, gt_dict,)
        elif self.cfg.MODEL.HEAD.TYPE == 'CENTER':
            loss, status = self.compute_losses_CENTER(pred_dict, gt_dict,)
        if return_status:
            return loss, status
        else:
            return loss
    def compute_losses_CENTER(self, pred_dict, gt_dict, return_status=True):
        # gt gaussian map
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)
        gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1)

        # Get boxes
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0,
                                                                                                           max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # compute location loss
        if 'score_map' in pred_dict:
            location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps)
        else:
            location_loss = torch.tensor(0.0, device=l1_loss.device)
        # weighted sum
        giou_loss = giou_loss.mean()
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "Loss/location": location_loss.item(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss
    def compute_losses_GFL_New(self, pred_dict, gt_dict):
        # GT
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_bbox_xyxy = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        gt_box_vec_xyxy = gt_bbox_xyxy.unsqueeze(1).repeat(1,self.search_feat_size*self.search_feat_size,1).view(-1,4)

        # Predictions
        pred_boxes = pred_dict['pred_boxes']
        pred_boxes = pred_boxes.view(self.bs, -1, 4)
        pred_logits = pred_dict['score_map'] # bs, H, W, class
        best_boxes = pred_dict['box_best'] 
        
        pred_logits_vec = pred_logits.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        # Fetch the box with max pred score. 
        pred_box_final_xyxy = box_cxcywh_to_xyxy(pred_boxes)
        pred_logits_flatten = pred_logits.permute(0, 2, 3,1).flatten(1)
        pred_box_pos = torch.gather(pred_box_final_xyxy,1,pred_logits_flatten.max(1, keepdim=True)[1][...,None].expand(-1,-1,4)).view(-1,4)

        # LOC Loss
        pred_giou, pred_iou = self.objective['giou'](pred_box_pos, gt_bbox_xyxy)  # (BN,4) (BN,4)
        l1 = self.objective['l1'](pred_box_pos, gt_bbox_xyxy, reduction='none')  # (BN,4) (BN,4)
        giou_loss = pred_giou.mean()
        l1_loss = l1.mean()

        # CLASS Loss
        pred_boxes_vec_xyxy = pred_box_final_xyxy.view(-1, 4)
        iou, _ = box_iou(pred_boxes_vec_xyxy, gt_box_vec_xyxy)
        gt_inbox_map =  generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        qfloss = self.objective['qfl'](pred_logits_vec, gt_inbox_map.reshape(-1), iou.detach())
        ce_loss = qfloss.mean()

        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + ce_loss\
            # + df_loss * 0.25\
            # + pool_loss * 0.25
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = pred_iou.mean().detach()
        mean_giou_loss = giou_loss.detach()
        ce_loss_ = ce_loss.detach()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss.item(),
                    # "Loss/dfl": df_loss.item(),
                    "Loss/qfl": ce_loss_.item(),
                    "IoU": mean_iou.item(),
                    }
        return loss, status



    def compute_losses_GFL(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)

        # Get pred
        pred_boxes = pred_dict['pred_boxes']
        # pred_center_score_unsig = pred_dict['score_map_c']
        # qs = pred_dict['qs']
        pred_quality_score = pred_dict['score_map']
        best_boxes = pred_dict['box_best']
        # attn = pred_dict['attn'] # bs, nhead, (81+256), (81+256)
        # corr2GT = pred_dict['corr2GT']
        

        q_score = pred_quality_score.permute(0, 2, 3,
                                      1).reshape(-1, 1)
                                           
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        
        gt_inbox_map = generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)
        gt_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        zero_map = gt_inbox_map.new_zeros(gt_inbox_map.flatten(1).shape)

        filtered_pred_init, filtered_gt_init, box_pos_index = filter_predictions_by_iou(gt_bbox, pred_boxes.view(self.bs, -1, 4), iou_threshold=0.7, k=1, maxk=100)
        filtered_box_map = zero_map.clone().flatten(0).scatter_(0, box_pos_index.cuda(), 1)
        gt_inbox_map_box = gt_inbox_map #* filtered_box_map.reshape(gt_inbox_map.shape)

        if self.cfg.TRAIN.TRAIN_CLS:
            gt_inbox_map = gt_inbox_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            gt_center_map = gt_center_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            # topk_map_for_box = topk_map_for_box * gt_dict['label'].view(-1)[:,None,None].repeat(1,
            #                     self.search_feat_size,self.search_feat_size)
            # topk_map_for_cls = topk_map_for_cls * gt_dict['label'].view(-1)[:,None,None].repeat(1,
            #                     self.search_feat_size,self.search_feat_size)
        
        # Duplicate GT box
        pos_box_ind = torch.nonzero(gt_inbox_map_box.flatten(1))
        gt_bboxes_multiple = gt_bbox[pos_box_ind[:,0],:]
        pred_boxes_slected = pred_boxes.reshape(self.bs, -1, 4)[pos_box_ind[:,0],pos_box_ind[:,1],:]
        pred_ltrb = pred_dict['ltrb_dis'] # already inside the GT box
        pred_ltrb_pos = pred_ltrb.reshape(self.bs, -1, 68)[pos_box_ind[:,0],pos_box_ind[:,1],:].reshape(-1, 17)

        gt_boxes_vec = box_xywh_to_xyxy(gt_bboxes_multiple).clamp(min=0.0,max=1.0)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes_slected).view(-1, 4)

        iou, _ = box_iou(pred_boxes_vec, gt_boxes_vec)
        iou = iou.detach()
        q_score_pos = pred_quality_score.squeeze(1).flatten(1)[pos_box_ind[:,0], pos_box_ind[:,1]] # BUG done
        

        # # Adaptive select the pos sample
        # pos_ind, selected_iou, select_mask = filter_scores(iou, inside_ind, bs=bs) # ([bs_ind, flatten_map_ind])
        # q_score_pos_unsig = q_score_unsig.flatten(1)[pos_ind[:,0], pos_ind[:,0][1]]
        # pred_box_pos = pred_boxes_vec[select_mask]
        # target_box_pos = gt_boxes_vec[select_mask]
        
        if self.cfg.TRAIN.TRAIN_CLS:
            nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
            gt_bbox = gt_bbox[nonzero_indices[:,0],:]
            best_boxes = best_boxes[nonzero_indices[:,0],:]

        top_gt_boxes_vec = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0,max=1.0)
        top_pred_boxes_vec = best_boxes

        giou_loss, _ = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)
        # giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        top_giou, top_iou = self.objective['giou'](top_pred_boxes_vec, top_gt_boxes_vec)  # (BN,4) (BN,4)
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec, reduction='none')  # (BN,4) (BN,4)

        # Map loss
        iou_score = zero_map.clone().flatten(1) 
        iou_score[pos_box_ind[:,0], pos_box_ind[:, 1]] = iou

        # Only fetch selected samples.
        
        target_ltrb_pos =  bbox2distance(grid_center_2d(gt_inbox_map_box, self.search_feat_size),
            gt_boxes_vec * self.search_feat_size, max_dis=self.search_feat_size).reshape(-1)# dont norm to 1

        # The weight for localization, from predicted cls socre.
        target_weight = q_score_pos.detach()

        # Map for classification.
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_dict['search_anno'][0]), self.search_feat_size)
        iou_gaussian = iou_score * gaussian_center_map.flatten(1)
        # target_gaussian_weight = gaussian_center_map.flatten(1)[inside_ind[:,0],inside_ind[:,1]] * target_weight
        # gt_gaussian_maps[-1].unsqueeze(1)

        # Focal loss
        # qf_loss = self.objective['dfocal'](inverse_sigmoid(q_score), topk_map_for_cls.reshape(-1,1))
        # QFL handle all the points inside the box.
        # qs_loss = self.objective['qfl'](qs.reshape(-1, 1), gt_q_map.reshape(-1), iou_score.reshape(-1))
        qf_loss = self.objective['qfl'](inverse_sigmoid(q_score), gt_inbox_map.reshape(-1), iou_score.reshape(-1))
        qf_loss = qf_loss.reshape(self.bs,-1).mean(1).sum()
        # DFL only handle the selected points inside the box
        df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        giou_loss = (giou_loss * target_weight).sum()/self.bs # weighted GIoU
        l1_loss = (l1_loss.mean(1) * target_weight).sum()/self.bs # weighted L1
                
        # corr loss
        # pool_loss = self.objective['bce'](corr2GT.view(-1), gt_dict['label'].view(-1))

        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss\
            + self.loss_weight['l1'] * l1_loss\
            + qf_loss\
            + df_loss * 0.25\
            # + pool_loss * 0.25
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = top_iou.detach().mean()
        mean_giou_loss = top_giou.detach().mean()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss.item(),
                #   "Loss/ctr": qs_loss.item(),
                    "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    # "Loss/pool": pool_loss.item(),
                    "IoU": mean_iou.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    def compute_losses_REPPOINTS_(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gaussian_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        center_ind_batch, center_ind = torch.where(gaussian_center_map.flatten(1)>=1)
        # Get pred
        pred_box_init_vec = pred_dict['init_box']       # xyxy
        pred_box_refine_vec = pred_dict['refine_box']   # xyxy
        pred_cls_score = pred_dict['score_map']
        bs = pred_cls_score.shape[0]

        pred_box_init = pred_box_init_vec.reshape(bs,-1,4)
        pred_box_refine = pred_box_refine_vec.reshape(bs,-1,4)
        pred_cls_score = pred_cls_score.permute(0, 2, 3,
                                      1).reshape(-1, 1) # (n,1)
        num_queries = pred_box_refine.shape[1]
                                           
        if torch.isnan(pred_box_refine).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        # Calc IoU to select the positive samples for loc(refine).
        gt_boxes_vec_all = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4) --> (n,4)
        iou, _ = box_iou(pred_box_init_vec, gt_boxes_vec_all)
        iou = iou.detach()

        pos_inds = torch.where(iou > 0.5)[0]

        pred_box_refine_pos = pred_box_refine_vec[pos_inds]
        gt_box_pos = gt_boxes_vec_all[pos_inds]


        pred_box_init_center = pred_box_init[center_ind_batch, center_ind]
        pred_box_refine_center = pred_box_refine[center_ind_batch, center_ind]
        gt_boxes_vec_center = box_xywh_to_xyxy(gt_bbox).view(-1, 4).clamp(min=0.0, max=1.0)


        giou_loss_init_center, iou_init_center = self.objective['giou'](pred_box_init_center, gt_boxes_vec_center)
        giou_loss_refine_center, iou_refine_center = self.objective['giou'](pred_box_refine_center, gt_boxes_vec_center)
        giou_loss_refine_pos, _ = self.objective['giou'](pred_box_refine_pos, gt_box_pos)
        l1_loss_center = self.objective['l1'](pred_box_refine_center, gt_boxes_vec_center)  # (BN,4) (BN,4)

        # Map loss
        qf_loss = self.objective['qfl'](pred_cls_score, gaussian_center_map.reshape(-1), gaussian_center_map.reshape(-1))
        qf_loss = qf_loss.mean()
        # DFL only handle the selected points inside the box
        # df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        # giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
        giou_loss_init_center = giou_loss_init_center.mean()
        giou_loss_refine_center = giou_loss_refine_center.mean()
        giou_loss_refine_pos = giou_loss_refine_pos.mean()
                
        # weighted sum
        loss = giou_loss_init_center * 2.0\
        + giou_loss_refine_center * 2.0\
        + giou_loss_refine_pos * 2.0\
        + qf_loss\
        + l1_loss_center * 5.0
            # + df_loss * 0.25\
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = iou_refine_center.detach().mean()
        mean_iou_init = iou_init_center.detach().mean()
        mean_giou_loss = giou_loss_refine_center.detach()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss_center.item(),
                #   "Loss/ctr": qs_loss.item(),
                    # "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                    "IoU_init": mean_iou_init.item(),
                #   "IoU_T": mean_iou.item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    def compute_losses_REPPOINTS(self, pred_dict, gt_dict):
        # Get GT boxes
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gaussian_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        center_ind_batch, center_ind = torch.where(gaussian_center_map.flatten(1)>=1)
        # Get pred
        pred_box_init_vec = pred_dict['init_box']       # xyxy
        pred_box_refine_vec = pred_dict['refine_box']   # xyxy
        pred_cls_score = pred_dict['score_map']
        # anchor_box_vec = pred_dict['anchor_box']
        # guess_size = pred_dict['guess_size']
        bs = pred_cls_score.shape[0]

        pred_box_init = pred_box_init_vec.reshape(bs,-1,4)
        pred_box_refine = pred_box_refine_vec.reshape(bs,-1,4)
        # anchor_box = anchor_box_vec.reshape(bs,-1,4)
        pred_cls_score = pred_cls_score.permute(0, 2, 3,
                                      1).squeeze(-1) # (n,1)
        pred_best_ind = pred_cls_score.view(self.bs, -1).max(1, keepdim=True)[1]
        pred_cls_score_vec = pred_cls_score.reshape(-1,1)
        num_queries = pred_box_refine.shape[1]
                                           
        if torch.isnan(pred_box_refine).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        # Calc IoU to select the positive samples for loc(refine).
        gt_boxes_vec_all = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4) --> (n,4)
        iou, _ = box_iou(pred_box_init_vec, gt_boxes_vec_all)
        iou = iou.detach()

        pos_inds = torch.where(iou > 0.5)[0]

        pred_box_refine_pos = pred_box_refine_vec[pos_inds]
        gt_box_pos = gt_boxes_vec_all[pos_inds]


        pred_box_init_center = pred_box_init[center_ind_batch, center_ind]
        pred_box_refine_center = pred_box_refine[center_ind_batch, center_ind]
        # pred_box_init_center = torch.gather(pred_box_init,1, pred_best_ind[:,None].repeat(1,1,4)).squeeze(1)
        pred_box_best = torch.gather(pred_box_refine,1, pred_best_ind[:,None].repeat(1,1,4)).squeeze(1)

        # anchor_box_center = anchor_box[center_ind_batch, center_ind]
        gt_boxes_vec_center = box_xywh_to_xyxy(gt_bbox).view(-1, 4).clamp(min=0.0, max=1.0)


        if self.cfg.TRAIN.TRAIN_CLS:
            gt_q_map = gt_q_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
                                self.search_feat_size,self.search_feat_size)
            # gt_c_map = gt_c_map * gt_dict['label'].view(-1)[:,None,None].repeat(1,
            #                     self.search_feat_size,self.search_feat_size)
        

        
        if self.cfg.TRAIN.TRAIN_CLS:
            nonzero_indices=torch.nonzero(gt_dict['label'].view(-1))
            gt_bbox = gt_bbox[nonzero_indices[:,0],:]
            best_boxes = best_boxes[nonzero_indices[:,0],:]


        giou_loss_init_center, iou_init_center = self.objective['giou'](pred_box_init_center, gt_boxes_vec_center)
        giou_loss_refine_center, iou_refine_center = self.objective['giou'](pred_box_refine_center, gt_boxes_vec_center)
        _, iou_pred = self.objective['giou'](pred_box_best, gt_boxes_vec_center)
        l1_loss_init_center = self.objective['l1'](pred_box_init_center, gt_boxes_vec_center) 
        l1_loss_final_center = self.objective['l1'](pred_box_refine_center, gt_boxes_vec_center) 
        giou_loss_refine_pos, _ = self.objective['giou'](pred_box_refine_pos, gt_box_pos)
        l1_loss_refine_pos = self.objective['l1'](pred_box_refine_pos, gt_box_pos) 
         # (BN,4) (BN,4)
        # _, iou_anchor = self.objective['giou'](anchor_box_center, gt_boxes_vec_center)

        # a = box_xyxy_to_xywh(anchor_box_center)
        # b = box_xyxy_to_xywh(gt_boxes_vec_center)
        # Map loss
        qf_loss = self.objective['qfl'](pred_cls_score_vec, gaussian_center_map.reshape(-1), gaussian_center_map.reshape(-1))
        # qf_loss = self.objective['focal'](pred_cls_score_vec.sigmoid(), gaussian_center_map.reshape(-1,1))
        qf_loss = qf_loss.reshape(self.bs,-1).mean()
        # DFL only handle the selected points inside the box
        # df_loss = self.objective['dfl'](pred_ltrb_pos, target_ltrb_pos, weight=target_weight[:, None].expand(-1, 4).reshape(-1))
        # giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
        giou_loss_init_center = giou_loss_init_center.mean()
        giou_loss_refine_center = giou_loss_refine_center.mean()
        giou_loss_refine_pos = giou_loss_refine_pos.mean()
                
        # weighted sum
        loss = giou_loss_init_center * 2\
        + giou_loss_refine_center * 2\
        + qf_loss\
        + l1_loss_init_center * 5\
        + l1_loss_final_center * 5\
        + giou_loss_refine_pos  * 2\
        + l1_loss_refine_pos
            # + df_loss * 0.25\
            # + qs_loss * 0.25
            # + ctr_loss * 0 \
        
        # status for log
        mean_iou = iou_pred.detach().mean()
        mean_iou_init = iou_init_center.detach().mean()
        mean_giou_loss = giou_loss_refine_center.detach()
        
        status = {"Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/l1": l1_loss_final_center.item(),
                #   "Loss/ctr": qs_loss.item(),
                    # "Loss/dfl": df_loss.item(),
                    "Loss/qfl": qf_loss.item(),
                    "IoU": mean_iou.item(),
                    "IoU_init": mean_iou_init.item(),
                    # "IoU_anchor": iou_anchor.mean().item(),
                #   "Coord":pred_boxes_vec[0].detach().cpu()
                    }
        return loss, status
    
    def compute_losses_DECODER(self, pred_dict, gt_dict):
        # K = 60
        # num_classes = 1
        no_init_ce = False
        iou_aware_ce = True
        outputs_without_aux = {
            k: v for k, v in pred_dict.items() if k != "aux_outputs" and k != "enc_outputs"
        }
        # Get init(enc) and final
        pred_query_ind = outputs_without_aux['query_index']
        # pred_query_attn = outputs_without_aux['query_attn']
        src_logits = outputs_without_aux['pred_logits']
        _, K, num_classes = src_logits.shape
        pred_box_init_vec = outputs_without_aux['init_pred_boxes'].view(-1,4)      # xywh
        pred_box_final_vec = outputs_without_aux['pred_boxes'].reshape(-1,4)      # xywh
        pred_logit_init = outputs_without_aux['init_pred_logits']
        pred_logit_final = outputs_without_aux['pred_logits']
        pred_init_all_logits = outputs_without_aux['all_init_logits']
        pred_logit_init_vec = pred_logit_init.reshape(-1, num_classes)
        pred_logit_final_vec = pred_logit_final.reshape(-1, num_classes) 
        # all_init_logits = outputs_without_aux['all_init_logits'].view(self.bs, self.search_feat_size, self.search_feat_size)


        # GT
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_box_vec = gt_bbox.unsqueeze(1).repeat(1,K,1).view(-1,4)

        gt_box_vec_xyxy = box_xywh_to_xyxy(gt_box_vec).clamp(min=0.0, max=1.0)
        
        iou_matrix_init, _ = box_iou(gt_box_vec_xyxy, box_cxcywh_to_xyxy(pred_box_init_vec))
        iou_matrix_init_ = iou_matrix_init.reshape(self.bs, -1)
        iou_matrix_final, _ = box_iou(gt_box_vec_xyxy, box_cxcywh_to_xyxy(pred_box_final_vec))
        iou_matrix_final_ = iou_matrix_final.reshape(self.bs, -1)
        # filtered_pred_init, filtered_gt_init, box_pos_index_init = filter_predictions_by_iou(gt_bbox, outputs_without_aux['init_pred_boxes'], iou_threshold=0.7, maxk=2)
        # filtered_pred_init, filtered_gt_final, box_pos_index_final = filter_predictions_by_iou(gt_bbox, outputs_without_aux['init_pred_boxes'], iou_threshold=0.7, maxk=2)
        # _, _, cls_pos_index_init = filter_predictions_by_iou(gt_bbox, outputs_without_aux['init_pred_boxes'], iou_threshold=0.7, maxk=K, dynamic=False)
        _, _, cls_pos_index_final = filter_predictions_by_iou(gt_bbox, outputs_without_aux['init_pred_boxes'], iou_threshold=0.7, maxk=5, dynamic=False)
        # _, _, cls_nonneg_index = filter_predictions_by_iou(gt_bbox, outputs_without_aux['pred_boxes'], iou_threshold=0.5, maxk=K, dynamic=False)
        # box_pos_index = torch.arange()
        # CDN
        dn_meta = pred_dict['dn_meta']
        if dn_meta is not None and 'output_known_lbs_bboxes' in dn_meta:
            output_known_lbs_bboxes, single_padding, dn_num_group = prep_for_dn(dn_meta)
            scalar = dn_num_group
            dn_pred_box = output_known_lbs_bboxes['pred_boxes']
            dn_pred_logits = output_known_lbs_bboxes['pred_logits']
            dn_gt_box = gt_bbox.unsqueeze(1).repeat(1,dn_num_group*single_padding,1).view(-1,4)
            dn_iou, _ = box_iou(box_xywh_to_xyxy(dn_gt_box).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(dn_pred_box.reshape(-1,4)))
            # fetch the 1st box of each dn group
            dn_pred_box_positive = dn_pred_box.unsqueeze(1).reshape(self.bs, dn_num_group, single_padding, 4)[:, :, 0].reshape(-1,4)
            dn_gt_box_pos = gt_bbox.unsqueeze(1).repeat(1,dn_num_group,1).view(-1,4)
            dn_gt_labels = torch.linspace(1,0,2)[None].repeat(self.bs, dn_num_group).cuda()
            # iou_cdn = 
                                           
        if torch.isnan(pred_box_init_vec).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        
        query_map = torch.zeros([self.bs, self.search_feat_size * self.search_feat_size], device=pred_logit_final_vec.device)
        query_map_ = query_map.clone()
        query_map.scatter_(1, pred_query_ind, 1)
        # query_map_.scatter_(1, pred_query_attn, 1)
        query_map_ = query_map_.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        
        gt_box_map =  generate_cls_map(box_xywh_to_xyxy(gt_bbox), self.search_feat_size)

        gt_inbox_labels = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        gt_cls_labels_init = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        gt_cls_labels_final = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        new_gt_cls_labels_final = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape) -1
        # gt_cls_labels_final = pred_logit_final_vec.new_zeros(pred_logit_final_vec.shape)
        # gt_cls_labels = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.shape[2]+1],
        #                                     dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
        # gt_labels[box_pos_index] = 1
        new_gt_cls_labels_final.view(-1)[iou_matrix_final<0.5] = 0
        # new_gt_cls_labels_final[cls_pos_index_final] = 1
        supervision_idx = (new_gt_cls_labels_final>=0)
        new_gt_cls_labels_final = new_gt_cls_labels_final[supervision_idx]
        # gt_cls_labels_init[cls_pos_index_init] = 1
        # gt_cls_labels_final[cls_pos_index_final] = 1
        gt_cls_labels_final[cls_pos_index_final] = 1
        # pred_box_init_pos =  box_cxcywh_to_xyxy(pred_box_init_vec[box_pos_index]).clamp(min=0.0, max=1.0)
        # pred_box_final_pos =  box_cxcywh_to_xyxy(pred_box_final_vec[box_pos_index]).clamp(min=0.0, max=1.0)
        # gt_box_pos = box_xywh_to_xyxy(filtered_gt_init).clamp(min=0.0, max=1.0)

        

        # status for log
        scores = outputs_without_aux['pred_logits'].squeeze(dim=2) 
        top_scores, top_indices = torch.max(scores, dim=1)
        pred_top_ind = torch.gather(pred_query_ind,1,top_indices[:,None])
        pred_top_map = torch.zeros_like(query_map).scatter_(1, pred_top_ind, 1)
        iou_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, iou_matrix_final.reshape(self.bs, -1))
        init_iou_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, iou_matrix_init.reshape(self.bs, -1))
        pred_cls_map_init = torch.zeros_like(query_map, dtype=torch.float16).scatter_(1, pred_query_ind, pred_logit_init.squeeze(-1))
        pred_cls_map_final = torch.zeros_like(query_map, dtype=torch.float16).scatter_(1, pred_query_ind, pred_logit_final.squeeze(-1).sigmoid())
        # pred_attn_map = torch.zeros_like(query_map).scatter_(1, pred_query_ind, pred_query_attn)
        # topk_iou_map = torch.zeros_like(query_map, dtype=torch.float16).scatter_(1, pred_query_ind, gt_cls_labels.reshape(self.bs, -1))
        # label_cross_map = pred_attn_map * gt_box_map.reshape(self.bs,-1)
        # targets = torch.gather(gt_box_map.flatten(1), 1, pred_query_ind)

        # DEBUG maps
        pred_top_map = pred_top_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        query_map = query_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        iou_map = iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        init_iou_map = init_iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        # topk_iou_map = topk_iou_map.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        pred_cls_map_init = pred_cls_map_init.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        pred_cls_map_final = pred_cls_map_final.reshape(self.bs, self.search_feat_size, self.search_feat_size)
        gaussian_center_map = get_2d_gaussian_map(xywh_to_cxcywh(gt_bbox), self.search_feat_size, alpha=0.5)
        norm_center_map = generate_heatmap(gt_dict['search_anno'], self.search_feat_size)[0]
        inbox_topk_map = query_map * gt_box_map
        box_targets = torch.gather(inbox_topk_map.flatten(1), 1, pred_query_ind)
        gaussian_center_targets = torch.gather(norm_center_map.flatten(1), 1, pred_query_ind)
        q_in_box_ind = box_targets.reshape(-1).nonzero().reshape(-1)

        num_cls_samples = reduce_mean(
            torch.tensor(box_targets.sum()).cuda()).item()
        # num_cls_samples = 1
        topK = 4
        num_box_samples = reduce_mean(
            torch.tensor(self.bs*topK, dtype=torch.float32).cuda()).item()
        if self.count == 12:
            print("###############num_box_samples##############")
            print(num_box_samples)
        gt_inbox_labels[q_in_box_ind] = 1
        
        inbox_gt_box = gt_box_vec_xyxy[gt_inbox_labels.squeeze(1) > 0]
        inbox_pred_box_init = box_cxcywh_to_xyxy(pred_box_init_vec[gt_inbox_labels.squeeze(1) > 0])
        inbox_pred_box_final = box_cxcywh_to_xyxy(pred_box_final_vec[gt_inbox_labels.squeeze(1) > 0])
        # topN_and_inbox_bit =  torch.isin(torch.where(gt_inbox_labels.squeeze(1) > 0)[0], box_pos_index_final.cuda())
        # inbox_pred_box_init = inbox_pred_box_init[topN_and_inbox_bit]
        # inbox_pred_box_final = inbox_pred_box_final[topN_and_inbox_bit]
        # inbox_gt_box = inbox_gt_box[topN_and_inbox_bit]
        # gt_cls_labels[cls_pos_index] = 1

        # Every first query as positive box
        
        pred_box_init_xyxy = box_cxcywh_to_xyxy(outputs_without_aux['init_pred_boxes'])
        pred_box_final_xyxy = box_cxcywh_to_xyxy(outputs_without_aux['pred_boxes'])
        pred_box_final_pos = torch.gather(pred_box_final_xyxy,1,pred_logit_final[...,0].topk(topK)[1][...,None].expand(-1,-1,4)).view(-1,4)
        pred_box_init_pos = torch.gather(pred_box_init_xyxy,1,pred_logit_final[...,0].topk(topK)[1][...,None].expand(-1,-1,4)).view(-1,4)
        gt_bbox_xyxy = box_xywh_to_xyxy(gt_bbox).clamp(min=0.0, max=1.0)
        gt_box_pos = gt_box_init_pos = gt_bbox_xyxy[:,None].repeat(1,topK,1).view(-1, 4)

        # gt_box_pos = box_xywh_to_xyxy(filtered_gt_final).clamp(min=0.0, max=1.0)
        # gt_box_init_pos = box_xywh_to_xyxy(filtered_gt_init).clamp(min=0.0, max=1.0)
        # pred_box_init_pos =  box_cxcywh_to_xyxy(pred_box_init_vec[box_pos_index_init])
        # pred_box_final_pos =  box_cxcywh_to_xyxy(pred_box_final_vec[box_pos_index_final])
        # gt_box_pos = box_xywh_to_xyxy(filtered_gt_init).clamp(min=0.0, max=1.0)

        ############################
        #  Simple Hungarian Match  #
        ############################
        # box_giou_cost = self.objective['giou'](gt_box_vec_xyxy, box_cxcywh_to_xyxy(pred_box_final_vec))[0]-1 # bs,N
        # box_l1_cost = torch.pairwise_distance(gt_box_vec_xyxy, box_cxcywh_to_xyxy(pred_box_final_vec), p=1)
        # out_prob = pred_logit_final_vec.sigmoid()
        # alpha = 0.25
        # gamma = 2.0
        # neg_cost_class = (1 - alpha) * (out_prob ** gamma) * (-(1 - out_prob + 1e-5).log())
        # pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-5).log())
        # cost_class = pos_cost_class[:, 0] - neg_cost_class[:, 0]
        # Cost = 2.0*cost_class + 5.0*box_l1_cost + 2.0*box_giou_cost
        # min_ind = Cost.view(self.bs,-1).min(1,keepdim=True)[1] # bs, 1
        # cls_targets = src_logits.new_zeros(src_logits.shape).squeeze(-1)
        # cls_targets.scatter_(1, min_ind, 1)
        ############################
        #           End            #
        ############################
        
        # pred_box_init_pos = torch.gather(pred_box_init_xyxy,1, min_ind[...,None].expand(-1,-1,4)).view(-1,4)
        # pred_box_final_pos = torch.gather(pred_box_final_xyxy,1, min_ind[...,None].expand(-1,-1,4)).view(-1,4)


        
        # ce_loss_init = self.objective['qfl'](inverse_sigmoid(pred_logit_init_vec), gt_box_map.reshape(-1), iou_matrix_init.detach())
        # LOSS
        if iou_aware_ce:
            # ce_loss_init = self.objective['dfocal'](inverse_sigmoid(pred_logit_init_vec) , iou_matrix_init[:,None].detach())
            if no_init_ce:
                ce_loss_init = 0
            else:
                ce_loss_init = self.objective['qfl'](pred_logit_init_vec, box_targets.reshape(-1), iou_matrix_init.detach())
            # ce_loss_final = self.objective['qfl'](pred_logit_final_vec, box_targets.reshape(-1), iou_matrix_final.detach())
            # ce_loss_final = self.objective['focal'](pred_logit_final_vec.sigmoid().view(self.bs,-1), gaussian_center_targets)
            ce_loss_final = self.objective['qfl'](pred_logit_final_vec, gt_inbox_labels.reshape(-1), iou_matrix_init.detach())
            # ce_loss_init = self.objective['qfl'](inverse_sigmoid(pred_logit_init_vec), gt_inbox_labels.reshape(-1), iou_matrix_init.detach())
            # ce_loss_final = self.objective['dfocal'](pred_logit_final_vec, gt_cls_labels_final)
        else:
            # ce_loss_init = self.objective['dfocal'](pred_logit_init, gt_cls_labels_init.reshape(self.bs,-1, num_classes))
            # ce_loss_final = self.objective['dfocal'](pred_logit_final, gt_cls_labels_final.reshape(self.bs,-1, num_classes))
            # ce_loss_init = self.objective['dfocal'](pred_logit_init_vec, gt_cls_labels_init)
            ce_loss_init = 0
            # pred_logit_final_vec = pred_logit_final_vec[supervision_idx]
            # ce_loss_final = self.objective['dfocal']((pred_logit_final_vec), gaussian_center_map.flatten(1)[...,None])
            # ce_loss_final = self.objective['dfocal']((pred_logit_final_vec), cls_targets.view(-1,1))
            # ce_loss_init = self.objective['qfl'](pred_logit_init_vec, box_targets.reshape(-1), iou_matrix_init.detach())
            ce_loss_final = self.objective['qfl'](pred_logit_final_vec, box_targets.reshape(-1), iou_matrix_final.detach())
            # ce_loss_final = self.objective['qfl'](pred_logit_final_vec, gt_inbox_labels.reshape(-1), iou_matrix_final.detach())
            # ce_loss_init = self.objective['focal'](pred_init_all_logits.sigmoid().view(self.bs,-1), gaussian_center_targets)
            # ce_loss_final = self.objective['focal'](pred_logit_final_vec.sigmoid().view(self.bs,-1), gaussian_center_targets)
        
        # ce_loss_final = self.objective['qfl'](pred_logit_final_vec, box_targets.reshape(-1), iou_matrix_final)
        giou_loss_init, iou_init = self.objective['giou'](pred_box_init_pos, gt_box_init_pos)
        giou_loss_final, iou_final = self.objective['giou'](pred_box_final_pos, gt_box_pos)
        l1_loss_init = self.objective['l1'](pred_box_init_pos, gt_box_init_pos)  # (BN,4) (BN,4)
        l1_loss_final = self.objective['l1'](pred_box_final_pos, gt_box_pos)  # (BN,4) (BN,4)
        # giou_loss_init, iou_init = self.objective['giou'](inbox_pred_box_init, inbox_gt_box)
        # giou_loss_final, iou_final = self.objective['giou'](inbox_pred_box_final, inbox_gt_box)
        # l1_loss_init = self.objective['l1'](inbox_pred_box_init, inbox_gt_box)  # (BN,4) (BN,4)
        # l1_loss_final = self.objective['l1'](inbox_pred_box_final, inbox_gt_box)  # (BN,4) (BN,4)

        # CDN LOSS
        if dn_meta is not None:
            dn_pred_box_pos =  box_cxcywh_to_xyxy(dn_pred_box_positive)
            dn_gt_box_pos = box_xywh_to_xyxy(dn_gt_box_pos).clamp(min=0.0, max=1.0)
            giou_loss_cdn, iou_cdn_pos = self.objective['giou'](dn_pred_box_pos, dn_gt_box_pos)
            if iou_aware_ce:
                # ce_loss_cdn = self.objective['dfocal'](dn_pred_logits.reshape(-1,1), dn_iou.view(-1,1).detach())
                # ce_loss_cdn = self.objective['focal'](dn_pred_logits.sigmoid().view(self.bs,-1), dn_gt_labels)
                ce_loss_cdn = self.objective['qfl'](dn_pred_logits.reshape(-1,1), dn_gt_labels.reshape(-1), dn_iou.detach())
            else:
                ce_loss_cdn = self.objective['qfl'](dn_pred_logits.reshape(-1,1), dn_gt_labels.view(-1), dn_iou.detach())
                # ce_loss_cdn = self.objective['dfocal'](dn_pred_logits.reshape(-1,1), dn_gt_labels.view(-1,1))
            # ce_loss_cdn = ce_loss_cdn.mean(1).reshape(self.bs,dn_num_group*single_padding).mean(1).sum()
            ce_loss_cdn = ce_loss_cdn.sum() / num_cls_samples / scalar
            
            l1_loss_cdn = self.objective['l1'](dn_pred_box_pos, dn_gt_box_pos, reduction='none')  # (BN,4) (BN,4)
            loss_cdn = ce_loss_cdn * 1.0 + 2.0 * giou_loss_cdn.sum() / num_box_samples / scalar + 5.0 * l1_loss_cdn.sum() / num_box_samples / scalar
            aux_loss_cdn = 0
            for aux_output in output_known_lbs_bboxes['aux_outputs']:
                _dn_pred_box = aux_output['pred_boxes']
                _boxes_vec = _dn_pred_box.unsqueeze(1).reshape(self.bs, dn_num_group, single_padding, 4)[:, :, 0].reshape(-1,4)
                _logits_vec = aux_output['pred_logits'].reshape(-1, num_classes)
                _boxes_pos = box_cxcywh_to_xyxy(_boxes_vec)
                _giou_loss, _iou = self.objective['giou'](_boxes_pos, dn_gt_box_pos)
                _l1_loss = self.objective['l1'](_boxes_pos, dn_gt_box_pos, reduction='none')  # (BN,4) (BN,4)
                _dn_iou, _ = box_iou(box_xywh_to_xyxy(dn_gt_box).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(_dn_pred_box.reshape(-1,4)))
                if iou_aware_ce:
                    _ce_loss = self.objective['qfl'](_logits_vec.reshape(-1,1), dn_gt_labels.reshape(-1), _dn_iou.detach())
                    # _ce_loss = self.objective['focal'](_logits_vec.sigmoid().view(self.bs,-1), dn_gt_labels)
                    # _ce_loss = self.objective['dfocal'](_logits_vec,  _dn_iou.view(-1,1).detach())
                else:
                    _ce_loss = self.objective['qfl'](_logits_vec.reshape(-1,1), dn_gt_labels.reshape(-1), _dn_iou.detach())
                    # _ce_loss = self.objective['dfocal'](_logits_vec, dn_gt_labels.view(-1,1))
                # _ce_loss_cdn = _ce_loss.mean(1).reshape(self.bs,dn_num_group*single_padding).mean(1).sum()
                _ce_loss = _ce_loss.sum() / num_cls_samples / scalar
                _l1_loss = self.objective['l1'](_boxes_pos, dn_gt_box_pos)  # (BN,4) (BN,4)
                aux_loss_cdn += 1.0 * _ce_loss + 2.0 * _giou_loss.sum() / num_box_samples / scalar + 5.0 *_l1_loss.sum() / num_box_samples / scalar
        else:
            loss_cdn = torch.as_tensor(0.0).to("cuda")
            aux_loss_cdn = torch.as_tensor(0.0).to("cuda")

        # AUX LOSS
        aux_loss = 0
        for aux_output in pred_dict['aux_outputs']:
            _logits = aux_output['pred_logits']
            _logits_vec = _logits.reshape(-1, num_classes)
            _boxes_vec = aux_output['pred_boxes']#.reshape(-1,4)
            # _boxes_pos = box_cxcywh_to_xyxy(_boxes_vec[box_pos_index_final])
            _boxes_pos = torch.gather(_boxes_vec,1,_logits[...,0].topk(topK)[1][...,None].expand(-1,-1,4)).view(-1,4)
            _boxes_pos = box_cxcywh_to_xyxy(_boxes_pos)
            _giou_loss, _ = self.objective['giou'](_boxes_pos, gt_box_pos)
            _iou, _ = box_iou(box_xywh_to_xyxy(gt_box_vec).clamp(min=0.0, max=1.0), box_cxcywh_to_xyxy(_boxes_vec.reshape(-1,4)))
            if iou_aware_ce:
                # _ce_loss = self.objective['dfocal'](inverse_sigmoid(_logits_vec) , _iou[:,None].detach())
                _ce_loss = self.objective['qfl'](_logits_vec, box_targets.reshape(-1), _iou.detach())
                # _ce_loss = self.objective['focal'](_logits_vec.sigmoid().view(self.bs,-1), gaussian_center_targets)
                # _ce_loss = self.objective['dfocal'](_logits_vec, gt_cls_labels_final)
                # _ce_loss = self.objective['qfl'](_logits_vec, gt_inbox_labels.reshape(-1), _iou.detach())
            else:
                # _ce_loss = self.objective['dfocal'](_logits_vec, cls_targets.view(-1,1))
                _ce_loss = self.objective['qfl'](_logits_vec, box_targets.reshape(-1), _iou.detach())
                # _ce_loss = self.objective['dfocal'](_logits, gt_cls_labels_final.reshape(self.bs,-1, num_classes))
                # _ce_loss = self.objective['dfocal'](_logits_vec[supervision_idx,], new_gt_cls_labels_final)
                # _ce_loss = self.objective['focal'](_logits_vec.sigmoid().view(self.bs,-1), gaussian_center_targets)
            # _ce_loss = _ce_loss.mean(1).reshape(self.bs,K).sum(1).mean()
            _ce_loss = _ce_loss.sum() / num_cls_samples
            # _ce_loss = 0
            _l1_loss = self.objective['l1'](_boxes_pos, gt_box_pos, reduction='none')  # (BN,4) (BN,4)
            aux_loss += 1.0 * _ce_loss + 2.0 * _giou_loss.sum() / num_box_samples + 5.0 *_l1_loss.sum() / num_box_samples

        # giou_loss = (giou_loss * target_weight).mean() # weighted GIoU
        # ce_loss_init = ce_loss_init.mean(1).reshape(self.bs,-1).sum(1).mean()
        # ce_loss_final = ce_loss_final.mean(1).reshape(self.bs,K).sum(1).mean()
        ce_loss_final = ce_loss_final.sum() / num_cls_samples
        if not no_init_ce:
            ce_loss_init = ce_loss_init.sum() / num_cls_samples #.reshape(self.bs,K).mean(1).sum()
        # ce_loss_init = ce_loss_init / K
        giou_loss_init = giou_loss_init.sum() / num_box_samples
        giou_loss_final = giou_loss_final.sum() / num_box_samples
        # save_masked_img(gt_dict, 11,alpha=0.15,boxxywh=outputs_without_aux['pred_boxes'])
        init_box_loss = giou_loss_init * 2.0 + l1_loss_init * 5.0
        final_box_loss = giou_loss_final * 2.0 + l1_loss_final * 5.0
                
        # weighted sum
        loss = 0 +\
        + ce_loss_init * 1.0\
        + ce_loss_final * 1.0\
        + init_box_loss * 1.0\
        + final_box_loss * 1.0\
        + aux_loss * 1.0\
        + loss_cdn * 1.0\
        + aux_loss_cdn * 1.0

        
        
        
        # targets = gt_cls_labels.new_zeros((self.bs, 256).shape)
        # for i, m in enumerate(label_cross_map):
        #     if m.sum()<=0:
        #         max_id = gaussian_center_map[i].view(-1).argmax()
        #     else:
        #         max_id = m.view(-1).argmax()
        #     targets[i, max_id] = 1

        top_boxes = torch.gather(outputs_without_aux['pred_boxes'], 1, top_indices[:,None,None].expand(-1, -1, 4)).squeeze(dim=1) 
        pred_iou, _ = box_iou(box_cxcywh_to_xyxy(top_boxes), box_xywh_to_xyxy(gt_bbox).clamp(min=0.0, max=1.0))
        query_corre_num = (gt_box_map*query_map).sum(1).sum(1)
        # search_img = gt_dict['search_images'][0][7]*torch.asarray([0.229, 0.224, 0.225], device=scores.device)[:,None,None]+\
        #     torch.asarray([0.485, 0.456, 0.406], device=scores.device)[:,None,None]
        # torchvision.utils.save_image(search_img, 'xxxxxx.png')
        # F.interpolate(gt_box_map[None,:], size=(256,256))
        # torchvision.utils.save_image((gt_dict['search_images'][0][2]*torch.asarray([0.229, 0.224, 0.225], device=scores.device)[:,None,None]+\
        #    torch.asarray([0.485, 0.456, 0.406], device=scores.device)[:,None,None])*F.interpolate(gt_box_map[None,:], size=(256,256))[0,2], 'xxxxxx.png')
        mean_iou_pred = pred_iou.detach().mean()
        # mean_iou_attn = iou_matrix_final_[:,0].mean()
        mean_iou_init = iou_init.detach().mean()
        mean_iou_last = iou_final.detach().mean()
        cls_loss = ce_loss_final.detach()
        mean_giou_loss = giou_loss_final.detach()
        # avg_box_pos_num = len(box_pos_index_init)//self.bs
        # avg_cls_pos_num = len(cls_pos_index_final)//self.bs

        status = {  "Loss/total": loss.item(),
                    "Loss/giou": mean_giou_loss.item(),
                    "Loss/cls": cls_loss.item(),
                    "IoU_i": mean_iou_init.item(),
                    "IoU_f": mean_iou_last.item(),
                    # "IoU_a": mean_iou_attn.item(),
                    "IoU": mean_iou_pred.item(),
                    # "Pb":avg_box_pos_num,
                    # "Pc":avg_cls_pos_num
                    }
        if len(pred_dict['aux_outputs'])>0:
            status['Loss/aux'] = aux_loss.item()
        return loss, status
     
# [save_masked_img(gt_dict, i, pred_cls_map_final, xyxy=pred_box_final_pos.view(self.bs,-1,4), f_name=f'xxx{i}.png') for i in range(self.bs)]
# save_masked_img(gt_dict, 0, pred_cls_map_final, xyxy=box_xywh_to_xyxy(gt_bbox[:,None]))
# save_masked_img(gt_dict, 0, pred_cls_map_final, xyxy=pred_box_final_xyxy)
# [save_masked_img(gt_dict, i, pred_cls_map_final, xyxy=box_xywh_to_xyxy(gt_bbox[:,None]), f_name=f'xxx{i}.png') for i in range(self.bs)]
def save_masked_img(gt_dict, batch_id, mask_map=None, src='search', f_name='xxx.png', alpha=.2, xyxy=None):
    image_batch = gt_dict[src+'_images'][0,batch_id].cpu()
    std = torch.asarray([0.229, 0.224, 0.225])[:,None,None]
    mean = torch.asarray([0.485, 0.456, 0.406])[:,None,None]

    shape = gt_dict[src+'_images'].shape[-1]

    if mask_map is not None:
        mask = F.interpolate(mask_map[None,:], size=shape)[0, batch_id].cpu() + alpha
    else:
        mask = 1
    img = (image_batch * std + mean)*mask
    if xyxy is not None:
        img_int = (img*256).clamp(0,255.9).to(torch.uint8)
        boxxyxy = xyxy[batch_id] * shape
        img_int = torchvision.utils.draw_bounding_boxes(img_int, boxxyxy, width=1)
        img = img_int.to(torch.float32)/256.0
    torchvision.utils.save_image(img, f_name)


def filter_scores(S, Ind, bs, alpha=1.0):
    # Compute the batch indices
    batch_indices = Ind[:, 0]

    # Compute the mean and standard error for each batch
    batch_means = torch.zeros(bs, dtype=torch.float32, device=S.device)
    batch_std = torch.zeros(bs, dtype=torch.float32, device=S.device)
    for i in range(bs):
        mask = (Ind[:, 0] == i)
        batch_scores = S[mask]
        batch_size = batch_scores.size(0)
        if batch_size > 0:
            batch_means[i] = batch_scores.mean()
            batch_std[i] = batch_scores.std() / torch.sqrt(torch.tensor(batch_size))

    # Compute the threshold for each batch
    batch_thresholds = batch_means + alpha * batch_std

    # Compute the mask for the scores that pass the threshold
    # mask = S > batch_thresholds[batch_indices]
    mask = S > 0

    # Filter the indices and scores
    filtered_indices = Ind[mask]
    filtered_scores = S[mask]

    return filtered_indices, filtered_scores, mask

   
def batch_box_iou(boxes1, boxes2):
    area1 = boxes1[:, 2] * boxes1[:, 3]
    area2 = boxes2[:, 2] * boxes2[:, 3]

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [bs, N, 2]
    rb = torch.min(boxes1[:, None, 2:4] + boxes1[:, None, :2], boxes2[:, 2:4] + boxes2[:, :2])  # [bs, N, 2]

    wh = (rb - lt).clamp(min=0)  # [bs, N, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [bs, N]

    iou = inter / (area1[:, None] + area2 - inter)
    return iou

def filter_predictions_by_iou(GT_boxes, pred_boxes, iou_threshold=0.5, k=1, maxk=5, dynamic=False):
    bs, n = pred_boxes.shape[:2]

    # Calculate IoU between GT_boxes and pred_boxes
    GT_vec = box_xywh_to_xyxy(GT_boxes)[:, None, :].repeat((1, n, 1)).view(-1, 4).clamp(min=0.0, max=1.0)
    P_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4).clamp(min=0.0, max=1.0)
    iou, _ = box_iou(GT_vec, P_vec)
    iou = iou.view(bs,n)

    # Filter predictions based on IoU threshold
    filtered_preds = []
    filtered_gts = []
    positive_indices = []
    for i in range(bs):
        if dynamic:
            iou_threshold = iou[i].mean()+iou[i].std()
        valid_indices = (iou[i] >= iou_threshold).nonzero(as_tuple=True)[0]
        
        if len(valid_indices) >= maxk:
            topk_indices = iou[i].topk(maxk, largest=True, sorted=True).indices
            valid_indices = topk_indices
        # If no predictions meet the condition, reserve the top-k predictions
        if len(valid_indices) == 0:
            topk_indices = iou[i].topk(k, largest=True, sorted=True).indices
            valid_indices = topk_indices

        filtered_preds.append(pred_boxes[i, valid_indices])
        filtered_gts.append(GT_boxes[i].expand(len(valid_indices), -1))

        # Store the indices of positive samples
        positive_indices.extend((i * n + valid_indices).tolist())

    return torch.cat(filtered_preds), torch.cat(filtered_gts), torch.tensor(positive_indices)

def __filter_predictions_by_iou(gt_boxes, pred_boxes, iou_threshold=0.5, k=1):
    bs, n = pred_boxes.shape[:2]

    # Function to compute IoU between two boxes
    def iou(box1, box2):
        x1, y1, x2, y2 = box1
        X1, Y1, X2, Y2 = box2

        w_intersection = max(0, min(x2, X2) - max(x1, X1))
        h_intersection = max(0, min(y2, Y2) - max(y1, Y1))
        area_intersection = w_intersection * h_intersection

        area_box1 = (x2 - x1) * (y2 - y1)
        area_box2 = (X2 - X1) * (Y2 - Y1)
        area_union = area_box1 + area_box2 - area_intersection

        return area_intersection / area_union

    filtered_pred_boxes = []
    filtered_gt_boxes = []

    for i in range(bs):
        gt_box = gt_boxes[i]
        preds = pred_boxes[i]

        # Compute IoU values between the GT box and predicted boxes
        iou_values = torch.tensor([iou(gt_box, pred) for pred in preds])

        # Filter predictions by the IoU threshold
        filtered_preds = preds[iou_values >= iou_threshold]

        # If no predictions meet the threshold, keep the top-k predictions based on the IoU values
        if len(filtered_preds) == 0:
            topk_indices = iou_values.topk(k, dim=0)[1]
            filtered_preds = preds[topk_indices]

        filtered_pred_boxes.append(filtered_preds)
        filtered_gt_boxes.append(gt_box.unsqueeze(0).repeat(len(filtered_preds),1))

    return filtered_pred_boxes, filtered_gt_boxes

def prep_for_dn(dn_meta):
        output_known_lbs_bboxes = dn_meta['output_known_lbs_bboxes']
        num_dn_groups,pad_size=dn_meta['num_dn_group'],dn_meta['pad_size']
        assert pad_size % num_dn_groups==0
        single_pad=pad_size//num_dn_groups

        return output_known_lbs_bboxes,single_pad,num_dn_groups
lib/train/actors/ostrack.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy
import torch
from lib.utils.merge import merge_template_search


class STARKSActor(BaseActor):
    """ Actor for training the STARK-S and STARK-ST(Stage1)"""
    def __init__(self, net, objective, loss_weight, settings):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # forward pass
        out_dict = self.forward_pass(data, run_box_head=True, run_cls_head=False)
        # print(out_dict.keys())
        # process the groundtruth
        gt_bboxes = data['search_anno']  # (Ns, batch, 4) (x1,y1,w,h)

        # compute losses
        loss, status = self.compute_losses(out_dict, gt_bboxes[0])

        return loss, status

    def forward_pass(self, data, run_box_head, run_cls_head):
        feat_dict_list = []
        # process the templates
        # print("-------------------------HELLO",data.keys())
        # raise Exception("caonima")
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            feat_dict_list.append(self.net(img=NestedTensor(template_img_i, template_att_i), mode='backbone'))

        # process the search regions (t-th frame)
        search_img = data['search_images'].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        search_att = data['search_att'].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)
        feat_dict_list.append(self.net(img=NestedTensor(search_img, search_att), mode='backbone'))

        # run the transformer and compute losses
        seq_dict = merge_template_search(feat_dict_list)
        out_dict, _, _ = self.net(seq_dict=seq_dict, mode="transformer", run_box_head=run_box_head, run_cls_head=run_cls_head)
        # out_dict: (B, N, C), outputs_coord: (1, B, N, C), target_query: (1, B, N, C)
        return out_dict

    def compute_losses(self, pred_dict, gt_bbox, return_status=True):
        # Get boxes
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss

lib/train/actors/stark_s.py
TsingWei/LiteTrack
from . import BaseActor
from lib.utils.misc import NestedTensor
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy
import torch
from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate


class TokenTrackActor(BaseActor):
    """ Actor for training OSTrack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        assert len(data['search_images']) == 1

        template_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1,
                                                             *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        box_mask_z = None
        ce_keep_rate = None
        if self.cfg.MODEL.BACKBONE.CE_LOC:
            box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
                                            data['template_anno'][0])

            ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
            ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
            ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
                                                total_epochs=ce_start_epoch + ce_warm_epoch,
                                                ITERS_PER_EPOCH=1,
                                                base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])
        template_bb = None
        if self.cfg.MODEL.USE_CLS_TOKEN:
            template_bb = data['template_anno'][0]

        if len(template_list) == 1:
            template_list = template_list[0]

        out_dict = self.net(template=template_list,
                            search=search_img,
                            ce_template_mask=box_mask_z,
                            ce_keep_rate=ce_keep_rate,
                            return_last_attn=False,
                            template_bb=template_bb)

        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        # gt gaussian map
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE)
        gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1)

        # Get boxes
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0,
                                                                                                           max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()
        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # compute location loss
        if 'score_map' in pred_dict:
            location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps)
        else:
            location_loss = torch.tensor(0.0, device=l1_loss.device)
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "Loss/location": location_loss.item(),
                      "IoU": mean_iou.item(),
                    #   "Coord":pred_boxes_vec[0].detach().cpu()
                      }
            return loss, status
        else:
            return loss

lib/train/actors/tokentrack.py
TsingWei/LiteTrack
import importlib
import os
from collections import OrderedDict


def create_default_local_file():
    path = os.path.join(os.path.dirname(__file__), 'local.py')

    empty_str = '\'\''
    default_settings = OrderedDict({
        'workspace_dir': empty_str,
        'tensorboard_dir': 'self.workspace_dir + \'/tensorboard/\'',
        'pretrained_networks': 'self.workspace_dir + \'/pretrained_networks/\'',
        'lasot_dir': empty_str,
        'got10k_dir': empty_str,
        'trackingnet_dir': empty_str,
        'coco_dir': empty_str,
        'lvis_dir': empty_str,
        'sbd_dir': empty_str,
        'imagenet_dir': empty_str,
        'imagenetdet_dir': empty_str,
        'ecssd_dir': empty_str,
        'hkuis_dir': empty_str,
        'msra10k_dir': empty_str,
        'davis_dir': empty_str,
        'youtubevos_dir': empty_str})

    comment = {'workspace_dir': 'Base directory for saving network checkpoints.',
               'tensorboard_dir': 'Directory for tensorboard files.'}

    with open(path, 'w') as f:
        f.write('class EnvironmentSettings:\n')
        f.write('    def __init__(self):\n')

        for attr, attr_val in default_settings.items():
            comment_str = None
            if attr in comment:
                comment_str = comment[attr]
            if comment_str is None:
                f.write('        self.{} = {}\n'.format(attr, attr_val))
            else:
                f.write('        self.{} = {}    # {}\n'.format(attr, attr_val, comment_str))


def create_default_local_file_ITP_train(workspace_dir, data_dir):
    path = os.path.join(os.path.dirname(__file__), 'local.py')

    empty_str = '\'\''
    default_settings = OrderedDict({
        'workspace_dir': workspace_dir,
        'tensorboard_dir': os.path.join(workspace_dir, 'tensorboard'),    # Directory for tensorboard files.
        'pretrained_networks': os.path.join(workspace_dir, 'pretrained_networks'),
        'lasot_dir': os.path.join(data_dir, 'lasot'),
        'got10k_dir': os.path.join(data_dir, 'got10k/train'),
        'got10k_val_dir': os.path.join(data_dir, 'got10k/val'),
        'lasot_lmdb_dir': os.path.join(data_dir, 'lasot_lmdb'),
        'got10k_lmdb_dir': os.path.join(data_dir, 'got10k_lmdb'),
        'trackingnet_dir': os.path.join(data_dir, 'trackingnet'),
        'trackingnet_lmdb_dir': os.path.join(data_dir, 'trackingnet_lmdb'),
        'coco_dir': os.path.join(data_dir, 'coco'),
        'coco_lmdb_dir': os.path.join(data_dir, 'coco_lmdb'),
        'lvis_dir': empty_str,
        'sbd_dir': empty_str,
        'imagenet_dir': os.path.join(data_dir, 'vid'),
        'imagenet_lmdb_dir': os.path.join(data_dir, 'vid_lmdb'),
        'imagenetdet_dir': empty_str,
        'ecssd_dir': empty_str,
        'hkuis_dir': empty_str,
        'msra10k_dir': empty_str,
        'davis_dir': empty_str,
        'youtubevos_dir': empty_str})

    comment = {'workspace_dir': 'Base directory for saving network checkpoints.',
               'tensorboard_dir': 'Directory for tensorboard files.'}

    with open(path, 'w') as f:
        f.write('class EnvironmentSettings:\n')
        f.write('    def __init__(self):\n')

        for attr, attr_val in default_settings.items():
            comment_str = None
            if attr in comment:
                comment_str = comment[attr]
            if comment_str is None:
                if attr_val == empty_str:
                    f.write('        self.{} = {}\n'.format(attr, attr_val))
                else:
                    f.write('        self.{} = \'{}\'\n'.format(attr, attr_val))
            else:
                f.write('        self.{} = \'{}\'    # {}\n'.format(attr, attr_val, comment_str))


def env_settings():
    env_module_name = 'lib.train.admin.local'
    try:
        env_module = importlib.import_module(env_module_name)
        return env_module.EnvironmentSettings()
    except:
        env_file = os.path.join(os.path.dirname(__file__), 'local.py')

        create_default_local_file()
        raise RuntimeError('YOU HAVE NOT SETUP YOUR local.py!!!\n Go to "{}" and set all the paths you need. Then try to run again.'.format(env_file))

lib/train/admin/environment.py
TsingWei/LiteTrack
import torch
from torch.utils.data.distributed import DistributedSampler
# datasets related
from lib.train.dataset import Lasot, Got10k, MSCOCOSeq, ImagenetVID, TrackingNet, UAV
from lib.train.dataset import Lasot_lmdb, Got10k_lmdb, MSCOCOSeq_lmdb, ImagenetVID_lmdb, TrackingNet_lmdb
from lib.train.data import sampler, opencv_loader, jpeg4py_loader, processing, LTRLoader
import lib.train.data.transforms as tfm
from lib.utils.misc import is_main_process


def update_settings(settings, cfg):
    settings.print_interval = cfg.TRAIN.PRINT_INTERVAL
    settings.search_area_factor = {'template': cfg.DATA.TEMPLATE.FACTOR,
                                   'search': cfg.DATA.SEARCH.FACTOR}
    settings.output_sz = {'template': cfg.DATA.TEMPLATE.SIZE,
                          'search': cfg.DATA.SEARCH.SIZE}
    settings.center_jitter_factor = {'template': cfg.DATA.TEMPLATE.CENTER_JITTER,
                                     'search': cfg.DATA.SEARCH.CENTER_JITTER}
    settings.scale_jitter_factor = {'template': cfg.DATA.TEMPLATE.SCALE_JITTER,
                                    'search': cfg.DATA.SEARCH.SCALE_JITTER}
    settings.grad_clip_norm = cfg.TRAIN.GRAD_CLIP_NORM
    settings.print_stats = None
    settings.batchsize = cfg.TRAIN.BATCH_SIZE
    settings.scheduler_type = cfg.TRAIN.SCHEDULER.TYPE


def names2datasets(name_list: list, settings, image_loader):
    assert isinstance(name_list, list)
    datasets = []
    for name in name_list:
        assert name in ["LASOT", "GOT10K_vottrain", "GOT10K_votval", "GOT10K_train_full", "GOT10K_official_val",
                        "COCO17", "VID", "TRACKINGNET","antiUAV"]
        if name == "LASOT":
            if settings.use_lmdb:
                print("Building lasot dataset from lmdb")
                datasets.append(Lasot_lmdb(settings.env.lasot_lmdb_dir, split='train', image_loader=image_loader))
            else:
                datasets.append(Lasot(settings.env.lasot_dir, split='train', image_loader=image_loader))
        if name == "GOT10K_vottrain":
            if settings.use_lmdb:
                print("Building got10k from lmdb")
                datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='vottrain', image_loader=image_loader))
            else:
                datasets.append(Got10k(settings.env.got10k_dir, split='vottrain', image_loader=image_loader))
        if name == "GOT10K_train_full":
            if settings.use_lmdb:
                print("Building got10k_train_full from lmdb")
                datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='train_full', image_loader=image_loader))
            else:
                datasets.append(Got10k(settings.env.got10k_dir, split='train_full', image_loader=image_loader))
        if name == "GOT10K_votval":
            if settings.use_lmdb:
                print("Building got10k from lmdb")
                datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='votval', image_loader=image_loader))
            else:
                datasets.append(Got10k(settings.env.got10k_dir, split='votval', image_loader=image_loader))
        if name == "GOT10K_official_val":
            if settings.use_lmdb:
                raise ValueError("Not implement")
            else:
                datasets.append(Got10k(settings.env.got10k_val_dir, split=None, image_loader=image_loader))
        if name == "COCO17":
            if settings.use_lmdb:
                print("Building COCO2017 from lmdb")
                datasets.append(MSCOCOSeq_lmdb(settings.env.coco_lmdb_dir, version="2017", image_loader=image_loader))
            else:
                datasets.append(MSCOCOSeq(settings.env.coco_dir, version="2017", image_loader=image_loader))
        if name == "VID":
            if settings.use_lmdb:
                print("Building VID from lmdb")
                datasets.append(ImagenetVID_lmdb(settings.env.imagenet_lmdb_dir, image_loader=image_loader))
            else:
                datasets.append(ImagenetVID(settings.env.imagenet_dir, image_loader=image_loader))
        if name == "TRACKINGNET":
            if settings.use_lmdb:
                print("Building TrackingNet from lmdb")
                datasets.append(TrackingNet_lmdb(settings.env.trackingnet_lmdb_dir, image_loader=image_loader))
            else:
                # raise ValueError("NOW WE CAN ONLY USE TRACKINGNET FROM LMDB")
                datasets.append(TrackingNet(settings.env.trackingnet_dir, image_loader=image_loader))
        if name == "antiUAV":
            if settings.use_lmdb:
                raise ValueError("Not support anti-UAV in lmdb format")
            else:
                datasets.append(UAV(settings.env.antiuav_dir, image_loader=image_loader))
    return datasets


def build_dataloaders(cfg, settings):
    # Data transform
    transform_joint = tfm.Transform(tfm.ToGrayscale(probability=0.05),
                                    tfm.RandomHorizontalFlip(probability=0.5)
                                    )

    transform_train_template = tfm.Transform(
                                    # tfm.RandomAffine(max_rotation=10),
                                    tfm.ToTensorAndJitter(0.2),
                                    # tfm.RandomHorizontalFlip_Norm(probability=0.5),
                                    tfm.Normalize(mean=cfg.DATA.MEAN, std=cfg.DATA.STD))
    transform_train_search = tfm.Transform(tfm.ToTensorAndJitter(0.2),
                                    # tfm.RandomHorizontalFlip_Norm(probability=0.5),
                                    tfm.Normalize(mean=cfg.DATA.MEAN, std=cfg.DATA.STD))

    transform_val = tfm.Transform(tfm.ToTensor(),
                                  tfm.Normalize(mean=cfg.DATA.MEAN, std=cfg.DATA.STD))

    # The tracking pairs processing module
    output_sz = settings.output_sz
    search_area_factor = settings.search_area_factor

    data_processing_train = processing.STARKProcessing(search_area_factor=search_area_factor,
                                                       output_sz=output_sz,
                                                       center_jitter_factor=settings.center_jitter_factor,
                                                       scale_jitter_factor=settings.scale_jitter_factor,
                                                       mode='sequence',
                                                       template_transform=transform_train_template,
                                                       search_transform=transform_train_search,
                                                       joint_transform=transform_joint,
                                                       settings=settings)

    data_processing_val = processing.STARKProcessing(search_area_factor=search_area_factor,
                                                     output_sz=output_sz,
                                                     center_jitter_factor=settings.center_jitter_factor,
                                                     scale_jitter_factor=settings.scale_jitter_factor,
                                                     mode='sequence',
                                                     transform=transform_val,
                                                     joint_transform=transform_joint,
                                                     settings=settings)

    # Train sampler and loader
    settings.num_template = getattr(cfg.DATA.TEMPLATE, "NUMBER", 1)
    settings.num_search = getattr(cfg.DATA.SEARCH, "NUMBER", 1)
    sampler_mode = getattr(cfg.DATA, "SAMPLER_MODE", "causal")
    train_cls = getattr(cfg.TRAIN, "TRAIN_CLS", False)
    pos_prob = getattr(cfg.TRAIN, "POSITIVE_PROB", 0.5)

    print("sampler_mode", sampler_mode)
    dataset_train = sampler.TrackingSampler(datasets=names2datasets(cfg.DATA.TRAIN.DATASETS_NAME, settings, jpeg4py_loader),
                                            p_datasets=cfg.DATA.TRAIN.DATASETS_RATIO,
                                            samples_per_epoch=cfg.DATA.TRAIN.SAMPLE_PER_EPOCH,
                                            max_gap=cfg.DATA.MAX_SAMPLE_INTERVAL, num_search_frames=settings.num_search,
                                            num_template_frames=settings.num_template, processing=data_processing_train,
                                            frame_sample_mode=sampler_mode, train_cls=train_cls, pos_prob=pos_prob)

    train_sampler = DistributedSampler(dataset_train) if settings.local_rank != -1 else None
    shuffle = False if settings.local_rank != -1 else True

    #DEBUG Time out here.
    loader_train = LTRLoader('train', dataset_train, training=True, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=shuffle,
                             num_workers=cfg.TRAIN.NUM_WORKER, drop_last=True, stack_dim=1, sampler=train_sampler, timeout=0)

    # Validation samplers and loaders
    dataset_val = sampler.TrackingSampler(datasets=names2datasets(cfg.DATA.VAL.DATASETS_NAME, settings, jpeg4py_loader),
                                          p_datasets=cfg.DATA.VAL.DATASETS_RATIO,
                                          samples_per_epoch=cfg.DATA.VAL.SAMPLE_PER_EPOCH,
                                          max_gap=cfg.DATA.MAX_SAMPLE_INTERVAL, num_search_frames=settings.num_search,
                                          num_template_frames=settings.num_template, processing=data_processing_val,
                                          frame_sample_mode=sampler_mode, train_cls=train_cls)
    val_sampler = DistributedSampler(dataset_val) if settings.local_rank != -1 else None
    loader_val = LTRLoader('val', dataset_val, training=False, batch_size=cfg.TRAIN.BATCH_SIZE,
                           num_workers=cfg.TRAIN.NUM_WORKER, drop_last=True, stack_dim=1, sampler=val_sampler,
                           epoch_interval=cfg.TRAIN.VAL_EPOCH_INTERVAL)

    return loader_train, loader_val


def get_optimizer_scheduler(net, cfg):
    train_cls = getattr(cfg.TRAIN, "TRAIN_CLS", False)
    if False:
        print("Only training classification head. Learnable parameters are shown below.")
        param_dicts = [
            {"params": [p for n, p in net.named_parameters() if "cls" in n and p.requires_grad]}
        ]

        for n, p in net.named_parameters():
            if "cls" not in n:
                p.requires_grad = False
            else:
                print(n)
    else:
        param_dicts = [
            {"params": [p for n, p in net.named_parameters() if "backbone" not in n and p.requires_grad]},
            {
                "params": [p for n, p in net.named_parameters() if "backbone" in n and p.requires_grad],
                "lr": cfg.TRAIN.LR * cfg.TRAIN.BACKBONE_MULTIPLIER,
            },
        ]
        # if is_main_process():
        #     print("Learnable parameters are shown below.")
        #     for n, p in net.named_parameters():
        #         if p.requires_grad:
        #             print(n)

    if cfg.TRAIN.OPTIMIZER == "ADAMW":
        optimizer = torch.optim.AdamW(param_dicts, lr=cfg.TRAIN.LR,
                                      weight_decay=cfg.TRAIN.WEIGHT_DECAY)
    else:
        raise ValueError("Unsupported Optimizer")
    if cfg.TRAIN.SCHEDULER.TYPE == 'step':
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.TRAIN.LR_DROP_EPOCH)
    elif cfg.TRAIN.SCHEDULER.TYPE == "Mstep":
        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                            milestones=cfg.TRAIN.SCHEDULER.MILESTONES,
                                                            gamma=cfg.TRAIN.SCHEDULER.GAMMA)
    else:
        raise ValueError("Unsupported scheduler")
    return optimizer, lr_scheduler

lib/train/base_functions.py
TsingWei/LiteTrack
from .lasot import Lasot
from .got10k import Got10k
from .tracking_net import TrackingNet
from .imagenetvid import ImagenetVID
from .coco import MSCOCO
from .coco_seq import MSCOCOSeq
from .got10k_lmdb import Got10k_lmdb
from .lasot_lmdb import Lasot_lmdb
from .imagenetvid_lmdb import ImagenetVID_lmdb
from .coco_seq_lmdb import MSCOCOSeq_lmdb
from .tracking_net_lmdb import TrackingNet_lmdb
from .uav import UAV

lib/train/dataset/__init__.py
TsingWei/LiteTrack
import os
import os.path
import numpy as np
import torch
import csv
import pandas
import random
from collections import OrderedDict
from .base_video_dataset import BaseVideoDataset
from lib.train.data import jpeg4py_loader
from lib.train.admin import env_settings


class Got10k(BaseVideoDataset):
    """ GOT-10k dataset.

    Publication:
        GOT-10k: A Large High-Diversity Benchmark for Generic Object Tracking in the Wild
        Lianghua Huang, Xin Zhao, and Kaiqi Huang
        arXiv:1810.11981, 2018
        https://arxiv.org/pdf/1810.11981.pdf

    Download dataset from http://got-10k.aitestunion.com/downloads
    """

    def __init__(self, root=None, image_loader=jpeg4py_loader, split=None, seq_ids=None, data_fraction=None):
        """
        args:
            root - path to the got-10k training data. Note: This should point to the 'train' folder inside GOT-10k
            image_loader (jpeg4py_loader) -  The function to read the images. jpeg4py (https://github.com/ajkxyz/jpeg4py)
                                            is used by default.
            split - 'train' or 'val'. Note: The validation split here is a subset of the official got-10k train split,
                    not NOT the official got-10k validation split. To use the official validation split, provide that as
                    the root folder instead.
            seq_ids - List containing the ids of the videos to be used for training. Note: Only one of 'split' or 'seq_ids'
                        options can be used at the same time.
            data_fraction - Fraction of dataset to be used. The complete dataset is used by default
        """
        root = env_settings().got10k_dir if root is None else root
        super().__init__('GOT10k', root, image_loader)

        # all folders inside the root
        self.sequence_list = self._get_sequence_list()

        # seq_id is the index of the folder inside the got10k root path
        if split is not None:
            if seq_ids is not None:
                raise ValueError('Cannot set both split_name and seq_ids.')
            ltr_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
            if split == 'train':
                file_path = os.path.join(ltr_path, 'data_specs', 'got10k_train_split.txt')
            elif split == 'val':
                file_path = os.path.join(ltr_path, 'data_specs', 'got10k_val_split.txt')
            elif split == 'train_full':
                file_path = os.path.join(ltr_path, 'data_specs', 'got10k_train_full_split.txt')
            elif split == 'vottrain':
                file_path = os.path.join(ltr_path, 'data_specs', 'got10k_vot_train_split.txt')
            elif split == 'votval':
                file_path = os.path.join(ltr_path, 'data_specs', 'got10k_vot_val_split.txt')
            else:
                raise ValueError('Unknown split name.')
            # seq_ids = pandas.read_csv(file_path, header=None, squeeze=True, dtype=np.int64).values.tolist()
            seq_ids = pandas.read_csv(file_path, header=None, dtype=np.int64).squeeze("columns").values.tolist()
        elif seq_ids is None:
            seq_ids = list(range(0, len(self.sequence_list)))

        self.sequence_list = [self.sequence_list[i] for i in seq_ids]

        if data_fraction is not None:
            self.sequence_list = random.sample(self.sequence_list, int(len(self.sequence_list)*data_fraction))

        self.sequence_meta_info = self._load_meta_info()
        self.seq_per_class = self._build_seq_per_class()

        self.class_list = list(self.seq_per_class.keys())
        self.class_list.sort()

    def get_name(self):
        return 'got10k'

    def has_class_info(self):
        return True

    def has_occlusion_info(self):
        return True

    def _load_meta_info(self):
        sequence_meta_info = {s: self._read_meta(os.path.join(self.root, s)) for s in self.sequence_list}
        return sequence_meta_info

    def _read_meta(self, seq_path):
        try:
            with open(os.path.join(seq_path, 'meta_info.ini')) as f:
                meta_info = f.readlines()
            object_meta = OrderedDict({'object_class_name': meta_info[5].split(': ')[-1][:-1],
                                       'motion_class': meta_info[6].split(': ')[-1][:-1],
                                       'major_class': meta_info[7].split(': ')[-1][:-1],
                                       'root_class': meta_info[8].split(': ')[-1][:-1],
                                       'motion_adverb': meta_info[9].split(': ')[-1][:-1]})
        except:
            object_meta = OrderedDict({'object_class_name': None,
                                       'motion_class': None,
                                       'major_class': None,
                                       'root_class': None,
                                       'motion_adverb': None})
        return object_meta

    def _build_seq_per_class(self):
        seq_per_class = {}

        for i, s in enumerate(self.sequence_list):
            object_class = self.sequence_meta_info[s]['object_class_name']
            if object_class in seq_per_class:
                seq_per_class[object_class].append(i)
            else:
                seq_per_class[object_class] = [i]

        return seq_per_class

    def get_sequences_in_class(self, class_name):
        return self.seq_per_class[class_name]

    def _get_sequence_list(self):
        with open(os.path.join(self.root, 'list.txt')) as f:
            dir_list = list(csv.reader(f))
        dir_list = [dir_name[0] for dir_name in dir_list]
        return dir_list

    def _read_bb_anno(self, seq_path):
        bb_anno_file = os.path.join(seq_path, "groundtruth.txt")
        gt = pandas.read_csv(bb_anno_file, delimiter=',', header=None, dtype=np.float32, na_filter=False, low_memory=False).values
        return torch.tensor(gt)

    def _read_target_visible(self, seq_path):
        # Read full occlusion and out_of_view
        occlusion_file = os.path.join(seq_path, "absence.label")
        cover_file = os.path.join(seq_path, "cover.label")

        with open(occlusion_file, 'r', newline='') as f:
            occlusion = torch.ByteTensor([int(v[0]) for v in csv.reader(f)])
        with open(cover_file, 'r', newline='') as f:
            cover = torch.ByteTensor([int(v[0]) for v in csv.reader(f)])

        target_visible = ~occlusion & (cover>0).byte()

        visible_ratio = cover.float() / 8
        return target_visible, visible_ratio

    def _get_sequence_path(self, seq_id):
        return os.path.join(self.root, self.sequence_list[seq_id])

    def get_sequence_info(self, seq_id):
        seq_path = self._get_sequence_path(seq_id)
        bbox = self._read_bb_anno(seq_path)

        valid = (bbox[:, 2] > 0) & (bbox[:, 3] > 0)
        visible, visible_ratio = self._read_target_visible(seq_path)
        visible = visible & valid.byte()

        return {'bbox': bbox, 'valid': valid, 'visible': visible, 'visible_ratio': visible_ratio}

    def _get_frame_path(self, seq_path, frame_id):
        return os.path.join(seq_path, '{:08}.jpg'.format(frame_id+1))    # frames start from 1

    def _get_frame(self, seq_path, frame_id):
        return self.image_loader(self._get_frame_path(seq_path, frame_id))

    def get_class_name(self, seq_id):
        obj_meta = self.sequence_meta_info[self.sequence_list[seq_id]]

        return obj_meta['object_class_name']

    def get_frames(self, seq_id, frame_ids, anno=None):
        seq_path = self._get_sequence_path(seq_id)
        obj_meta = self.sequence_meta_info[self.sequence_list[seq_id]]

        frame_list = [self._get_frame(seq_path, f_id) for f_id in frame_ids]

        if anno is None:
            anno = self.get_sequence_info(seq_id)

        anno_frames = {}
        for key, value in anno.items():
            anno_frames[key] = [value[f_id, ...].clone() for f_id in frame_ids]

        return frame_list, anno_frames, obj_meta

lib/train/dataset/got10k.py
TsingWei/LiteTrack
import os
import os.path
import torch
import numpy as np
import pandas
import csv
import random
from collections import OrderedDict
from .base_video_dataset import BaseVideoDataset
from lib.train.data import jpeg4py_loader
from lib.train.admin import env_settings


class Lasot(BaseVideoDataset):
    """ LaSOT dataset.

    Publication:
        LaSOT: A High-quality Benchmark for Large-scale Single Object Tracking
        Heng Fan, Liting Lin, Fan Yang, Peng Chu, Ge Deng, Sijia Yu, Hexin Bai, Yong Xu, Chunyuan Liao and Haibin Ling
        CVPR, 2019
        https://arxiv.org/pdf/1809.07845.pdf

    Download the dataset from https://cis.temple.edu/lasot/download.html
    """

    def __init__(self, root=None, image_loader=jpeg4py_loader, vid_ids=None, split=None, data_fraction=None):
        """
        args:
            root - path to the lasot dataset.
            image_loader (jpeg4py_loader) -  The function to read the images. jpeg4py (https://github.com/ajkxyz/jpeg4py)
                                            is used by default.
            vid_ids - List containing the ids of the videos (1 - 20) used for training. If vid_ids = [1, 3, 5], then the
                    videos with subscripts -1, -3, and -5 from each class will be used for training.
            split - If split='train', the official train split (protocol-II) is used for training. Note: Only one of
                    vid_ids or split option can be used at a time.
            data_fraction - Fraction of dataset to be used. The complete dataset is used by default
        """
        root = env_settings().lasot_dir if root is None else root
        super().__init__('LaSOT', root, image_loader)

        # Keep a list of all classes
        self.class_list = [f for f in os.listdir(self.root)]
        self.class_to_id = {cls_name: cls_id for cls_id, cls_name in enumerate(self.class_list)}

        self.sequence_list = self._build_sequence_list(vid_ids, split)

        if data_fraction is not None:
            self.sequence_list = random.sample(self.sequence_list, int(len(self.sequence_list)*data_fraction))

        self.seq_per_class = self._build_class_list()

    def _build_sequence_list(self, vid_ids=None, split=None):
        if split is not None:
            if vid_ids is not None:
                raise ValueError('Cannot set both split_name and vid_ids.')
            ltr_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
            if split == 'train':
                file_path = os.path.join(ltr_path, 'data_specs', 'lasot_train_split.txt')
            else:
                raise ValueError('Unknown split name.')
            # sequence_list = pandas.read_csv(file_path, header=None, squeeze=True).values.tolist()
            sequence_list = pandas.read_csv(file_path, header=None).squeeze("columns").values.tolist()
        elif vid_ids is not None:
            sequence_list = [c+'-'+str(v) for c in self.class_list for v in vid_ids]
        else:
            raise ValueError('Set either split_name or vid_ids.')

        return sequence_list

    def _build_class_list(self):
        seq_per_class = {}
        for seq_id, seq_name in enumerate(self.sequence_list):
            class_name = seq_name.split('-')[0]
            if class_name in seq_per_class:
                seq_per_class[class_name].append(seq_id)
            else:
                seq_per_class[class_name] = [seq_id]

        return seq_per_class

    def get_name(self):
        return 'lasot'

    def has_class_info(self):
        return True

    def has_occlusion_info(self):
        return True

    def get_num_sequences(self):
        return len(self.sequence_list)

    def get_num_classes(self):
        return len(self.class_list)

    def get_sequences_in_class(self, class_name):
        return self.seq_per_class[class_name]

    def _read_bb_anno(self, seq_path):
        bb_anno_file = os.path.join(seq_path, "groundtruth.txt")
        gt = pandas.read_csv(bb_anno_file, delimiter=',', header=None, dtype=np.float32, na_filter=False, low_memory=False).values
        return torch.tensor(gt)

    def _read_target_visible(self, seq_path):
        # Read full occlusion and out_of_view
        occlusion_file = os.path.join(seq_path, "full_occlusion.txt")
        out_of_view_file = os.path.join(seq_path, "out_of_view.txt")

        with open(occlusion_file, 'r', newline='') as f:
            occlusion = torch.ByteTensor([int(v) for v in list(csv.reader(f))[0]])
        with open(out_of_view_file, 'r') as f:
            out_of_view = torch.ByteTensor([int(v) for v in list(csv.reader(f))[0]])

        target_visible = ~occlusion & ~out_of_view

        return target_visible

    def _get_sequence_path(self, seq_id):
        seq_name = self.sequence_list[seq_id]
        class_name = seq_name.split('-')[0]
        vid_id = seq_name.split('-')[1]

        return os.path.join(self.root, class_name, class_name + '-' + vid_id)

    def get_sequence_info(self, seq_id):
        seq_path = self._get_sequence_path(seq_id)
        bbox = self._read_bb_anno(seq_path)

        valid = (bbox[:, 2] > 0) & (bbox[:, 3] > 0)
        visible = self._read_target_visible(seq_path) & valid.byte()

        return {'bbox': bbox, 'valid': valid, 'visible': visible}

    def _get_frame_path(self, seq_path, frame_id):
        return os.path.join(seq_path, 'img', '{:08}.jpg'.format(frame_id+1))    # frames start from 1

    def _get_frame(self, seq_path, frame_id):
        return self.image_loader(self._get_frame_path(seq_path, frame_id))

    def _get_class(self, seq_path):
        raw_class = seq_path.split('/')[-2]
        return raw_class

    def get_class_name(self, seq_id):
        seq_path = self._get_sequence_path(seq_id)
        obj_class = self._get_class(seq_path)

        return obj_class

    def get_frames(self, seq_id, frame_ids, anno=None):
        seq_path = self._get_sequence_path(seq_id)

        obj_class = self._get_class(seq_path)
        frame_list = [self._get_frame(seq_path, f_id) for f_id in frame_ids]

        if anno is None:
            anno = self.get_sequence_info(seq_id)

        anno_frames = {}
        for key, value in anno.items():
            anno_frames[key] = [value[f_id, ...].clone() for f_id in frame_ids]

        object_meta = OrderedDict({'object_class_name': obj_class,
                                   'motion_class': None,
                                   'major_class': None,
                                   'root_class': None,
                                   'motion_adverb': None})

        return frame_list, anno_frames, object_meta

lib/train/dataset/lasot.py
TsingWei/LiteTrack
import torch
import os
import os.path
import numpy as np
import csv
import pandas
from collections import OrderedDict

from lib.train.data import jpeg4py_loader
from lib.train.admin import env_settings
from .base_video_dataset import BaseVideoDataset


import glob
import json
import cv2

def list_sequences(root):

    sequence_list = []

    seq_dir = os.path.join(root)
    for filename in os.listdir(seq_dir):
        sequence_list.append(filename)

    return sequence_list


class UAV(BaseVideoDataset):

    def __init__(self, root=None, image_loader=jpeg4py_loader, split=None, seq_ids=None, data_fraction=None):

        root = env_settings().uav_dir if root is None else root
        super().__init__('antiUAV', root, image_loader)

        self.sequence_list = self._get_sequence_list()
        if seq_ids is None:
            seq_ids = list(range(0, len(self.sequence_list)))
        self.sequence_list = [self.sequence_list[i] for i in seq_ids]

        # self.sequence_list = list_sequences(self.root)
        self.ann_path = 'data/uav/train'

    def get_name(self):
        return 'uav'

    def get_num_sequences(self):
        return len(self.sequence_list)

    def _read_anno(self, anno_path):
        # gt = pandas.read_csv(anno_path, delimiter=',', header=None, dtype=np.float32, na_filter=False, low_memory=False).values
        res_file = os.path.join(anno_path, 'IR_label.json')
        with open(res_file, 'r') as f:
            label_res = json.load(f)
        gt = label_res['gt_rect']
        new_gt = [[0,0,0,0] if len(f)==0 else f for f in gt]
        try:
            a = torch.tensor(new_gt).float()
        except:
            print(anno_path)
        return torch.tensor(new_gt).float()

    def _read_target_visible(self, seq_path):
        res_file = os.path.join(seq_path, 'IR_label.json')
        with open(res_file, 'r') as f:
            label_res = json.load(f)
        exist = label_res['exist']
        target_visible = torch.ByteTensor([int(v) for v in exist])

        return target_visible

    def _get_sequence_path(self, seq_id):
        return os.path.join(self.root, self.sequence_list[seq_id])

    def _get_sequence_list(self):
        with open(os.path.join(self.root, 'list.txt')) as f:
            dir_list = list(csv.reader(f))
        dir_list = [dir_name[0] for dir_name in dir_list]
        return dir_list

    def get_sequence_info(self, seq_id):
        seq_path = self._get_sequence_path(seq_id)
        anno = self._read_anno(seq_path)
        valid = (anno[:,2]>0) & (anno[:,3]>0)
        visible = self._read_target_visible(seq_path)
        visible = visible & valid.byte()

        return {'bbox': anno, 'valid': valid, 'visible': visible}

    def _get_frame_path(self,seq_path,frame_id):
        # result = []
        # for filename in os.listdir(seq_path):
        #     result.append(int(filename[0:5]))

        return os.path.join(seq_path, '{:06}.jpg'.format(frame_id+1))

    def _get_frame(self, seq_path, frame_id):
        return self.image_loader(self._get_frame_path(seq_path, frame_id))
        # capture = cv2.VideoCapture(seq_path)
        # capture.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
        # ret, frame = capture.read()
        # import random
        # c = random.randint(0, 2)
        # kernel = np.ones((5, 5), np.uint8)
        # if random.random() < 0.5:
        #     return cv2.dilate(frame[:, :, c], kernel)
        # return frame[:, :, c]
        # return frame

    def get_frames(self, seq_id, frame_ids, anno=None):
        seq_path = self._get_sequence_path(seq_id)
        frame_list = [self._get_frame(seq_path, f) for f in frame_ids]

        if anno is None:
            anno = self.get_sequence_info(seq_id)

        # Create anno dict
        anno_frames = {}
        for key, value in anno.items():
            anno_frames[key] = [value[f_id, ...].clone() for f_id in frame_ids]

        object_meta = OrderedDict({'object_class': None,
                                   'motion_class': None,
                                   'major_class': None,
                                   'root_class': None,
                                   'motion_adverb': None})

        return frame_list, anno_frames, object_meta
lib/train/dataset/uav.py
TsingWei/LiteTrack
import os
import sys
import argparse
import importlib
import cv2 as cv
import torch.backends.cudnn
import torch.distributed as dist

import random
import numpy as np
torch.backends.cudnn.benchmark = False

import _init_paths
import lib.train.admin.settings as ws_settings


def init_seeds(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    # runner.set_random_seed(seed, deterministic=True)


def run_training(script_name, config_name, cudnn_benchmark=True, local_rank=-1, save_dir=None, base_seed=None,
                 use_lmdb=False, script_name_prv=None, config_name_prv=None, use_wandb=False,
                 distill=None, script_teacher=None, config_teacher=None):
    """Run the train script.
    args:
        script_name: Name of emperiment in the "experiments/" folder.
        config_name: Name of the yaml file in the "experiments/<script_name>".
        cudnn_benchmark: Use cudnn benchmark or not (default is True).
    """
    if save_dir is None:
        print("save_dir dir is not given. Use the default dir instead.")
    # This is needed to avoid strange crashes related to opencv
    cv.setNumThreads(0)
    cv.ocl.setUseOpenCL(False)

    torch.backends.cudnn.benchmark = cudnn_benchmark

    print('script_name: {}.py  config_name: {}.yaml'.format(script_name, config_name))

    '''2021.1.5 set seed for different process'''
    if base_seed is not None:
        if local_rank != -1:
            init_seeds(base_seed + local_rank)
        else:
            init_seeds(base_seed)

    settings = ws_settings.Settings()
    settings.script_name = script_name
    settings.config_name = config_name
    settings.project_path = 'train/{}/{}'.format(script_name, config_name)
    if script_name_prv is not None and config_name_prv is not None:
        settings.project_path_prv = 'train/{}/{}'.format(script_name_prv, config_name_prv)
    settings.local_rank = local_rank
    settings.save_dir = os.path.abspath(save_dir)
    settings.use_lmdb = use_lmdb
    prj_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
    settings.cfg_file = os.path.join(prj_dir, 'experiments/%s/%s.yaml' % (script_name, config_name))
    settings.use_wandb = use_wandb
    if distill:
        settings.distill = distill
        settings.script_teacher = script_teacher
        settings.config_teacher = config_teacher
        if script_teacher is not None and config_teacher is not None:
            settings.project_path_teacher = 'train/{}/{}'.format(script_teacher, config_teacher)
        settings.cfg_file_teacher = os.path.join(prj_dir, 'experiments/%s/%s.yaml' % (script_teacher, config_teacher))
        expr_module = importlib.import_module('lib.train.train_script_distill')
    else:
        expr_module = importlib.import_module('lib.train.train_script')
    expr_func = getattr(expr_module, 'run')

    expr_func(settings)


def main():
    parser = argparse.ArgumentParser(description='Run a train scripts in train_settings.')
    parser.add_argument('--script', type=str, required=True, help='Name of the train script.')
    parser.add_argument('--config', type=str, required=True, help="Name of the config file.")
    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Set cudnn benchmark on (1) or off (0) (default is on).')
    parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
    parser.add_argument('--save_dir', type=str, help='the directory to save checkpoints and logs')
    parser.add_argument('--seed', type=int, default=42, help='seed for random numbers')
    parser.add_argument('--use_lmdb', type=int, choices=[0, 1], default=0)  # whether datasets are in lmdb format
    parser.add_argument('--script_prv', type=str, default=None, help='Name of the train script of previous model.')
    parser.add_argument('--config_prv', type=str, default=None, help="Name of the config file of previous model.")
    parser.add_argument('--use_wandb', type=int, choices=[0, 1], default=0)  # whether to use wandb
    # for knowledge distillation
    parser.add_argument('--distill', type=int, choices=[0, 1], default=0)  # whether to use knowledge distillation
    parser.add_argument('--script_teacher', type=str, help='teacher script name')
    parser.add_argument('--config_teacher', type=str, help='teacher yaml configure file name')

    args = parser.parse_args()
    local_rank = int(os.environ.get("LOCAL_RANK",-1))
    if local_rank != -1:
        dist.init_process_group(backend='nccl')
        torch.cuda.set_device(local_rank)
    else:
        torch.cuda.set_device(0)
    run_training(args.script, args.config, cudnn_benchmark=args.cudnn_benchmark,
                 local_rank=local_rank, save_dir=args.save_dir, base_seed=args.seed,
                 use_lmdb=args.use_lmdb, script_name_prv=args.script_prv, config_name_prv=args.config_prv,
                 use_wandb=args.use_wandb,
                 distill=args.distill, script_teacher=args.script_teacher, config_teacher=args.config_teacher)


if __name__ == '__main__':
    main()

lib/train/run_training.py
TsingWei/LiteTrack
import os
# loss function related
from lib.utils.box_ops import giou_loss
from torch.nn.functional import l1_loss
from torch.nn import BCEWithLogitsLoss
# train pipeline related
from lib.train.trainers import LTRTrainer
# distributed training related
from torch.nn.parallel import DistributedDataParallel as DDP
# some more advanced functions
from .base_functions import *
# network related

from lib.models import build_ostrack, build_LiteTrack

from lib.train.actors import  OSTrackActor
# for import modules
import importlib
from lib.utils.slconfig import SLConfig
from ..utils.focal_loss import FocalLoss, DistributionFocalLoss, DFL, QFL, FL
import warnings
import torch.nn.functional as F
warnings.filterwarnings("ignore")

def run(settings):
    settings.description = 'Training script for STARK-S, STARK-ST stage1, and STARK-ST stage2'

    # update the default configs with config file
    if not os.path.exists(settings.cfg_file):
        raise ValueError("%s doesn't exist." % settings.cfg_file)
    config_module = importlib.import_module("lib.config.%s.config" % settings.script_name)
    cfg = config_module.cfg
    config_module.update_config_from_file(settings.cfg_file)
    if settings.local_rank in [-1, 0]:
        print("New configuration is shown below.")
        for key in cfg.keys():
            print("%s configuration:" % key, cfg[key])
            print('\n')

    # update settings based on cfg
    update_settings(settings, cfg)

    # Record the training log
    log_dir = os.path.join(settings.save_dir, 'logs')
    if settings.local_rank in [-1, 0]:
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
    settings.log_file = os.path.join(log_dir, "%s-%s.log" % (settings.script_name, settings.config_name))

    # Build dataloaders
    loader_train, loader_val = build_dataloaders(cfg, settings)


    # Create network
    if settings.script_name == "ostrack":
        net = build_ostrack(cfg)
    elif settings.script_name == "litetrack":
        net = build_LiteTrack(cfg)
    else:
        raise ValueError("illegal script name")

    # wrap networks to distributed one
    net.cuda()
    if settings.local_rank != -1:
        # net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)  # add syncBN converter
        net = DDP(net, device_ids=[settings.local_rank], find_unused_parameters=True)
        settings.device = torch.device("cuda:%d" % settings.local_rank)
    else:
        settings.device = torch.device("cuda:0")
    settings.deep_sup = getattr(cfg.TRAIN, "DEEP_SUPERVISION", False)
    settings.distill = getattr(cfg.TRAIN, "DISTILL", False)
    settings.distill_loss_type = getattr(cfg.TRAIN, "DISTILL_LOSS_TYPE", "KL")
    # Loss functions and Actors
    focal_loss = FocalLoss()
    distribution_focal_loss = FL()
    dfl = DFL()
    dfl_T = DFL()
    qfl = QFL()
    # negL = negL()
    if settings.script_name == "ostrack" or settings.script_name == "litetrack":
        objective = {'giou': giou_loss, 'l1': l1_loss, 'focal': focal_loss, 'dfocal':distribution_focal_loss,
                      'cls': BCEWithLogitsLoss(), 'dfl':dfl, 'qfl':qfl, 'dfl_T': dfl_T, 'bce': F.binary_cross_entropy_with_logits}
        loss_weight = {'giou': cfg.TRAIN.GIOU_WEIGHT, 'l1': cfg.TRAIN.L1_WEIGHT, 'focal': 1., 'cls': 1.0}
        actor = OSTrackActor(net=net, objective=objective, loss_weight=loss_weight, settings=settings, cfg=cfg)
    else:
        raise ValueError("illegal script name")

    # if cfg.TRAIN.DEEP_SUPERVISION:
    #     raise ValueError("Deep supervision is not supported now.")

    # Optimizer, parameters, and learning rates
    optimizer, lr_scheduler = get_optimizer_scheduler(net, cfg)
    use_amp = getattr(cfg.TRAIN, "AMP", False)
    trainer = LTRTrainer(actor, [loader_train, loader_val], optimizer, settings, lr_scheduler, use_amp=use_amp, cfg=cfg)

    # train process
    trainer.train(cfg.TRAIN.EPOCH, load_latest=True, fail_safe=True)

lib/train/train_script.py
TsingWei/LiteTrack
import os
# loss function related
from lib.utils.box_ops import giou_loss
from torch.nn.functional import l1_loss
from torch.nn import BCEWithLogitsLoss
# train pipeline related
from lib.train.trainers import LTRTrainer
# distributed training related
from torch.nn.parallel import DistributedDataParallel as DDP
# some more advanced functions
from .base_functions import *
# network related
from lib.models.stark import build_starks, build_starkst
from lib.models.stark import build_stark_lightning_x_trt
# forward propagation related
from lib.train.actors import STARKLightningXtrtdistillActor
# for import modules
import importlib


def build_network(script_name, cfg):
    # Create network
    if script_name == "stark_s":
        net = build_starks(cfg)
    elif script_name == "stark_st1" or script_name == "stark_st2":
        net = build_starkst(cfg)
    elif script_name == "stark_lightning_X_trt":
        net = build_stark_lightning_x_trt(cfg, phase="train")
    else:
        raise ValueError("illegal script name")
    return net


def run(settings):
    settings.description = 'Training script for STARK-S, STARK-ST stage1, and STARK-ST stage2'

    # update the default configs with config file
    if not os.path.exists(settings.cfg_file):
        raise ValueError("%s doesn't exist." % settings.cfg_file)
    config_module = importlib.import_module("lib.config.%s.config" % settings.script_name)
    cfg = config_module.cfg
    config_module.update_config_from_file(settings.cfg_file)
    if settings.local_rank in [-1, 0]:
        print("New configuration is shown below.")
        for key in cfg.keys():
            print("%s configuration:" % key, cfg[key])
            print('\n')

    # update the default teacher configs with teacher config file
    if not os.path.exists(settings.cfg_file_teacher):
        raise ValueError("%s doesn't exist." % settings.cfg_file_teacher)
    config_module_teacher = importlib.import_module("lib.config.%s.config" % settings.script_teacher)
    cfg_teacher = config_module_teacher.cfg
    config_module_teacher.update_config_from_file(settings.cfg_file_teacher)
    if settings.local_rank in [-1, 0]:
        print("New teacher configuration is shown below.")
        for key in cfg_teacher.keys():
            print("%s configuration:" % key, cfg_teacher[key])
            print('\n')

    # update settings based on cfg
    update_settings(settings, cfg)

    # Record the training log
    log_dir = os.path.join(settings.save_dir, 'logs')
    if settings.local_rank in [-1, 0]:
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
    settings.log_file = os.path.join(log_dir, "%s-%s.log" % (settings.script_name, settings.config_name))

    # Build dataloaders
    loader_train, loader_val = build_dataloaders(cfg, settings)

    if "RepVGG" in cfg.MODEL.BACKBONE.TYPE or "swin" in cfg.MODEL.BACKBONE.TYPE:
        cfg.ckpt_dir = settings.save_dir
    """turn on the distillation mode"""
    cfg.TRAIN.DISTILL = True
    cfg_teacher.TRAIN.DISTILL = True
    net = build_network(settings.script_name, cfg)
    net_teacher = build_network(settings.script_teacher, cfg_teacher)

    # wrap networks to distributed one
    net.cuda()
    net_teacher.cuda()
    net_teacher.eval()

    if settings.local_rank != -1:
        net = DDP(net, device_ids=[settings.local_rank], find_unused_parameters=True)
        net_teacher = DDP(net_teacher, device_ids=[settings.local_rank], find_unused_parameters=True)
        settings.device = torch.device("cuda:%d" % settings.local_rank)
    else:
        settings.device = torch.device("cuda:0")
    # settings.deep_sup = getattr(cfg.TRAIN, "DEEP_SUPERVISION", False)
    # settings.distill = getattr(cfg.TRAIN, "DISTILL", False)
    settings.distill_loss_type = getattr(cfg.TRAIN, "DISTILL_LOSS_TYPE", "L1")
    # Loss functions and Actors
    if settings.script_name == "stark_lightning_X_trt":
        objective = {'giou': giou_loss, 'l1': l1_loss}
        loss_weight = {'giou': cfg.TRAIN.GIOU_WEIGHT, 'l1': cfg.TRAIN.L1_WEIGHT}
        actor = STARKLightningXtrtdistillActor(net=net, objective=objective, loss_weight=loss_weight, settings=settings,
                                               net_teacher=net_teacher)
    else:
        raise ValueError("illegal script name")

    # Optimizer, parameters, and learning rates
    optimizer, lr_scheduler = get_optimizer_scheduler(net, cfg)
    use_amp = getattr(cfg.TRAIN, "AMP", False)
    trainer = LTRTrainer(actor, [loader_train, loader_val], optimizer, settings, lr_scheduler, use_amp=use_amp)

    # train process
    trainer.train(cfg.TRAIN.EPOCH, load_latest=True, fail_safe=True, distill=True)

lib/train/train_script_distill.py
TsingWei/LiteTrack
import os
import glob
import torch
import traceback
from lib.train.admin import multigpu
from torch.utils.data.distributed import DistributedSampler


class BaseTrainer:
    """Base trainer class. Contains functions for training and saving/loading checkpoints.
    Trainer classes should inherit from this one and overload the train_epoch function."""

    def __init__(self, actor, loaders, optimizer, settings, lr_scheduler=None):
        """
        args:
            actor - The actor for training the network
            loaders - list of dataset loaders, e.g. [train_loader, val_loader]. In each epoch, the trainer runs one
                        epoch for each loader.
            optimizer - The optimizer used for training, e.g. Adam
            settings - Training settings
            lr_scheduler - Learning rate scheduler
        """
        self.actor = actor
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.loaders = loaders

        self.update_settings(settings)

        self.epoch = 0
        self.stats = {}

        self.device = getattr(settings, 'device', None)
        if self.device is None:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() and settings.use_gpu else "cpu")

        self.actor.to(self.device)
        self.settings = settings

    def update_settings(self, settings=None):
        """Updates the trainer settings. Must be called to update internal settings."""
        if settings is not None:
            self.settings = settings

        if self.settings.env.workspace_dir is not None:
            self.settings.env.workspace_dir = os.path.expanduser(self.settings.env.workspace_dir)
            '''2021.1.4 New function: specify checkpoint dir'''
            if self.settings.save_dir is None:
                self._checkpoint_dir = os.path.join(self.settings.env.workspace_dir, 'checkpoints')
            else:
                self._checkpoint_dir = os.path.join(self.settings.save_dir, 'checkpoints')
            print("checkpoints will be saved to %s" % self._checkpoint_dir)

            if self.settings.local_rank in [-1, 0]:
                if not os.path.exists(self._checkpoint_dir):
                    print("Training with multiple GPUs. checkpoints directory doesn't exist. "
                          "Create checkpoints directory")
                    os.makedirs(self._checkpoint_dir)
        else:
            self._checkpoint_dir = None

    def train(self, max_epochs, load_latest=False, fail_safe=True, load_previous_ckpt=False, distill=False):
        """Do training for the given number of epochs.
        args:
            max_epochs - Max number of training epochs,
            load_latest - Bool indicating whether to resume from latest epoch.
            fail_safe - Bool indicating whether the training to automatically restart in case of any crashes.
        """

        epoch = -1
        num_tries = 1
        for i in range(num_tries):
            # try:
                if load_latest:
                    self.load_checkpoint()
                if load_previous_ckpt:
                    directory = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path_prv)
                    self.load_state_dict(directory)
                if distill:
                    directory_teacher = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path_teacher)
                    self.load_state_dict(directory_teacher, distill=True)
                for epoch in range(self.epoch+1, max_epochs+1):
                    self.epoch = epoch

                    self.train_epoch()

                    if self.lr_scheduler is not None:
                        if self.settings.scheduler_type != 'cosine':
                            self.lr_scheduler.step()
                        else:
                            self.lr_scheduler.step(epoch - 1)
                    # only save the last 10 checkpoints
                    save_every_epoch = getattr(self.settings, "save_every_epoch", False)
                    save_epochs = [79, 119, 120, 239]
                    if   epoch >= int(max_epochs*0.8)-1:
                        if self._checkpoint_dir:
                            if self.settings.local_rank in [-1, 0]:
                                self.save_checkpoint(eval=True)
                    elif epoch > (max_epochs - 1) or save_every_epoch or epoch % 40 == 0 or epoch in save_epochs:
                    # if epoch > (max_epochs - 10) or save_every_epoch or epoch % 100 == 0:
                        if self._checkpoint_dir:
                            if self.settings.local_rank in [-1, 0]:
                                self.save_checkpoint()
                                # self.delet_old_checkpoint(epoch-1)
                    
            # except:
            #     print('Training crashed at epoch {}'.format(epoch))
            #     if fail_safe:
            #         self.epoch -= 1
            #         load_latest = True
            #         print('Traceback for the error!')
            #         print(traceback.format_exc())
            #         print('Restarting training from last epoch ...')
            #     else:
            #         raise

        print('Finished training!')

    def train_epoch(self):
        raise NotImplementedError

    def save_checkpoint(self, eval=False):
        """Saves a checkpoint of the network and other variables."""

        net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net

        actor_type = type(self.actor).__name__
        net_type = type(net).__name__
        if eval:
            state = {
                'net': net.state_dict()
            }
        else:
            state = {
                'epoch': self.epoch,
                'actor_type': actor_type,
                'net_type': net_type,
                'net': net.state_dict(),
                'net_info': getattr(net, 'info', None),
                'constructor': getattr(net, 'constructor', None),
                'optimizer': self.optimizer.state_dict(),
                'stats': self.stats,
                'settings': self.settings
            }

        directory = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path)
        print(directory)
        if not os.path.exists(directory):
            print("directory doesn't exist. creating...")
            os.makedirs(directory)

        # First save as a tmp file
        tmp_file_path = '{}/{}_ep{:04d}.tmp'.format(directory, net_type, self.epoch)
        torch.save(state, tmp_file_path)

        file_path = '{}/{}_ep{:04d}.pth.tar'.format(directory, net_type, self.epoch)

        # Now rename to actual checkpoint. os.rename seems to be atomic if files are on same filesystem. Not 100% sure
        os.rename(tmp_file_path, file_path)

    def load_checkpoint(self, checkpoint = None, fields = None, ignore_fields = None, load_constructor = False):
        """Loads a network checkpoint file.

        Can be called in three different ways:
            load_checkpoint():
                Loads the latest epoch from the workspace. Use this to continue training.
            load_checkpoint(epoch_num):
                Loads the network at the given epoch number (int).
            load_checkpoint(path_to_checkpoint):
                Loads the file from the given absolute path (str).
        """

        net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net

        actor_type = type(self.actor).__name__
        net_type = type(net).__name__

        if checkpoint is None:
            # Load most recent checkpoint
            checkpoint_list = sorted(glob.glob('{}/{}/{}_ep*.pth.tar'.format(self._checkpoint_dir,
                                                                             self.settings.project_path, net_type)))
            if checkpoint_list:
                checkpoint_path = checkpoint_list[-1]
            else:
                print('No matching checkpoint file found')
                return
        elif isinstance(checkpoint, int):
            # Checkpoint is the epoch number
            checkpoint_path = '{}/{}/{}_ep{:04d}.pth.tar'.format(self._checkpoint_dir, self.settings.project_path,
                                                                 net_type, checkpoint)
        elif isinstance(checkpoint, str):
            # checkpoint is the path
            if os.path.isdir(checkpoint):
                checkpoint_list = sorted(glob.glob('{}/*_ep*.pth.tar'.format(checkpoint)))
                if checkpoint_list:
                    checkpoint_path = checkpoint_list[-1]
                else:
                    raise Exception('No checkpoint found')
            else:
                checkpoint_path = os.path.expanduser(checkpoint)
        else:
            raise TypeError

        # Load network
        checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')

        # assert net_type == checkpoint_dict['net_type'], 'Network is not of correct type.'

        if fields is None:
            fields = checkpoint_dict.keys()
        if ignore_fields is None:
            ignore_fields = ['settings']

            # Never load the scheduler. It exists in older checkpoints.
        ignore_fields.extend(['lr_scheduler', 'constructor', 'net_type', 'actor_type', 'net_info'])

        # Load all fields
        for key in fields:
            if key in ignore_fields:
                continue
            if key == 'net':
                net.load_state_dict(checkpoint_dict[key])
            elif key == 'optimizer':
                self.optimizer.load_state_dict(checkpoint_dict[key])
            else:
                setattr(self, key, checkpoint_dict[key])

        # Set the net info
        if load_constructor and 'constructor' in checkpoint_dict and checkpoint_dict['constructor'] is not None:
            net.constructor = checkpoint_dict['constructor']
        if 'net_info' in checkpoint_dict and checkpoint_dict['net_info'] is not None:
            net.info = checkpoint_dict['net_info']

        # Update the epoch in lr scheduler
        if 'epoch' in fields:
            self.lr_scheduler.last_epoch = self.epoch
        # 2021.1.10 Update the epoch in data_samplers
            for loader in self.loaders:
                if isinstance(loader.sampler, DistributedSampler):
                    loader.sampler.set_epoch(self.epoch)
        return True

    def delet_old_checkpoint(self, epoch_before=0):
        net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net
        net_type = type(net).__name__

        checkpoint_list = sorted(glob.glob('{}/{}/{}_ep*.pth.tar'.format(self._checkpoint_dir,
                                                                            self.settings.project_path, net_type)))
        
        checkpoint_list_to_remove = [ckpt_name for ckpt_name in  checkpoint_list if int(ckpt_name[-12:-8]) < epoch_before]
        for ckpt in checkpoint_list_to_remove:
            os.unlink(ckpt)
        # print(checkpoint_list_to_remove)
            

    def load_state_dict(self, checkpoint=None, distill=False):
        """Loads a network checkpoint file.

        Can be called in three different ways:
            load_checkpoint():
                Loads the latest epoch from the workspace. Use this to continue training.
            load_checkpoint(epoch_num):
                Loads the network at the given epoch number (int).
            load_checkpoint(path_to_checkpoint):
                Loads the file from the given absolute path (str).
        """
        if distill:
            net = self.actor.net_teacher.module if multigpu.is_multi_gpu(self.actor.net_teacher) \
                else self.actor.net_teacher
        else:
            net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net

        net_type = type(net).__name__

        if isinstance(checkpoint, str):
            # checkpoint is the path
            if os.path.isdir(checkpoint):
                checkpoint_list = sorted(glob.glob('{}/*_ep*.pth.tar'.format(checkpoint)))
                if checkpoint_list:
                    checkpoint_path = checkpoint_list[-1]
                else:
                    raise Exception('No checkpoint found')
            else:
                checkpoint_path = os.path.expanduser(checkpoint)
        else:
            raise TypeError

        # Load network
        print("Loading pretrained model from ", checkpoint_path)
        checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')

        assert net_type == checkpoint_dict['net_type'], 'Network is not of correct type.'

        missing_k, unexpected_k = net.load_state_dict(checkpoint_dict["net"], strict=False)
        print("previous checkpoint is loaded.")
        print("missing keys: ", missing_k)
        print("unexpected keys:", unexpected_k)

        return True

lib/train/trainers/base_trainer.py
TsingWei/LiteTrack
import os
import datetime
from collections import OrderedDict
import wandb
from lib.train.data.wandb_logger import WandbWriter
from lib.train.trainers import BaseTrainer
from lib.train.admin import AverageMeter, StatValue
from lib.train.admin import TensorboardWriter
import torch
import time
from torch.utils.data.distributed import DistributedSampler
from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler

from lib.utils.misc import get_world_size


class LTRTrainer(BaseTrainer):
    def __init__(self, actor, loaders, optimizer, settings, lr_scheduler=None, use_amp=False, cfg=None):
        """
        args:
            actor - The actor for training the network
            loaders - list of dataset loaders, e.g. [train_loader, val_loader]. In each epoch, the trainer runs one
                        epoch for each loader.
            optimizer - The optimizer used for training, e.g. Adam
            settings - Training settings
            lr_scheduler - Learning rate scheduler
        """
        super().__init__(actor, loaders, optimizer, settings, lr_scheduler)

        self._set_default_settings()

        # Initialize statistics variables
        self.stats = OrderedDict({loader.name: None for loader in self.loaders})

        # Initialize tensorboard and wandb
        self.wandb_writer = None
        if settings.local_rank in [-1, 0]:
            tensorboard_writer_dir = os.path.join(self.settings.env.tensorboard_dir, self.settings.project_path)
            if not os.path.exists(tensorboard_writer_dir):
                os.makedirs(tensorboard_writer_dir)
            self.tensorboard_writer = TensorboardWriter(tensorboard_writer_dir, [l.name for l in loaders])

            if settings.use_wandb:
                world_size = get_world_size()
                cur_train_samples = self.loaders[0].dataset.samples_per_epoch * max(0, self.epoch - 1)
                interval = (world_size * settings.batchsize)  # * interval
                self.wandb_writer = WandbWriter(settings.project_path[6:], cfg, tensorboard_writer_dir, cur_train_samples, interval,)

        self.move_data_to_gpu = getattr(settings, 'move_data_to_gpu', True)
        self.settings = settings
        self.use_amp = use_amp
        if use_amp:
            self.scaler = GradScaler()

    def _set_default_settings(self):
        # Dict of all default values
        default = {'print_interval': 10,
                   'print_stats': None,
                   'description': ''}

        for param, default_value in default.items():
            if getattr(self.settings, param, None) is None:
                setattr(self.settings, param, default_value)

    def cycle_dataset(self, loader):
        """Do a cycle of training or validation."""

        self.actor.train(loader.training)
        torch.set_grad_enabled(loader.training)

        self._init_timing()

        for i, data in enumerate(loader, 1):
            self.data_read_done_time = time.time()
            # get inputs
            if self.move_data_to_gpu:
                data = data.to(self.device)

            self.data_to_gpu_time = time.time()

            data['epoch'] = self.epoch
            data['settings'] = self.settings
            # forward pass
            if not self.use_amp:
                loss, stats = self.actor(data)
            else:
                with autocast():
                    loss, stats = self.actor(data)

            # backward pass and update weights
            if loader.training:
                self.optimizer.zero_grad()
                if not self.use_amp:
                    loss.backward()
                    if self.settings.grad_clip_norm > 0:
                        torch.nn.utils.clip_grad_norm_(self.actor.net.parameters(), self.settings.grad_clip_norm)
                    self.optimizer.step()
                else:
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()

            # update statistics
            batch_size = data['template_images'].shape[loader.stack_dim]
            self._update_stats(stats, batch_size, loader)

            # print statistics
            self._print_stats(i, loader, batch_size)

            # update wandb status
            if self.wandb_writer is not None and i % self.settings.print_interval == 0:

                # if self.settings.local_rank in [-1, 0]:
                #     x,x2,y,y2 = stats['Coord'][0],stats['Coord'][1],stats['Coord'][2],stats['Coord'][3]
                #     img = wandb.Image(data['search_images'][0][0],boxes={
                #         "predictions": {
                #             "box_data": [{
                #                 # one box expressed in the default relative/fractional domain
                #                 "position": {
                #                     "minX": x,
                #                     "maxX": x2,
                #                     "minY": y,
                #                     "maxY": y2
                #                     }
                #             }]}})
                            
                #     self.stats['train']['search_img'] = img
                    self.wandb_writer.write_log(self.stats, self.epoch)

        # calculate ETA after every epoch
        epoch_time = self.prev_time - self.start_time
        print("Epoch Time: " + str(datetime.timedelta(seconds=epoch_time)))
        print("Avg Data Time: %.5f" % (self.avg_date_time / self.num_frames * batch_size))
        print("Avg GPU Trans Time: %.5f" % (self.avg_gpu_trans_time / self.num_frames * batch_size))
        print("Avg Forward Time: %.5f" % (self.avg_forward_time / self.num_frames * batch_size))

    def train_epoch(self):
        """Do one epoch for each loader."""
        for loader in self.loaders:
            if self.epoch % loader.epoch_interval == 0:
                # 2021.1.10 Set epoch
                if isinstance(loader.sampler, DistributedSampler):
                    loader.sampler.set_epoch(self.epoch)
                self.cycle_dataset(loader)

        self._stats_new_epoch()
        if self.settings.local_rank in [-1, 0]:
            self._write_tensorboard()

    def _init_timing(self):
        self.num_frames = 0
        self.start_time = time.time()
        self.prev_time = self.start_time
        self.avg_date_time = 0
        self.avg_gpu_trans_time = 0
        self.avg_forward_time = 0

    def _update_stats(self, new_stats: OrderedDict, batch_size, loader):
        # Initialize stats if not initialized yet
        if loader.name not in self.stats.keys() or self.stats[loader.name] is None:
            self.stats[loader.name] = OrderedDict({name: AverageMeter() for name in new_stats.keys()})

        # add lr state
        if loader.training:
            lr_list = self.lr_scheduler.get_last_lr()
            for i, lr in enumerate(lr_list):
                var_name = 'LearningRate/group{}'.format(i)
                if var_name not in self.stats[loader.name].keys():
                    self.stats[loader.name][var_name] = StatValue()
                self.stats[loader.name][var_name].update(lr)

        for name, val in new_stats.items():
            if name not in self.stats[loader.name].keys():
                self.stats[loader.name][name] = AverageMeter()
            self.stats[loader.name][name].update(val, batch_size)

    def _print_stats(self, i, loader, batch_size):
        self.num_frames += batch_size
        current_time = time.time()
        batch_fps = batch_size / (current_time - self.prev_time)
        average_fps = self.num_frames / (current_time - self.start_time)
        prev_frame_time_backup = self.prev_time
        self.prev_time = current_time

        self.avg_date_time += (self.data_read_done_time - prev_frame_time_backup)
        self.avg_gpu_trans_time += (self.data_to_gpu_time - self.data_read_done_time)
        self.avg_forward_time += current_time - self.data_to_gpu_time

        if i % self.settings.print_interval == 0 or i == loader.__len__():
            print_str = '[%s: %d, %d / %d] ' % (loader.name, self.epoch, i, loader.__len__())
            print_str += 'FPS: %.1f (%.1f)  ,  ' % (average_fps, batch_fps)

            # 2021.12.14 add data time print
            print_str += 'DataTime: %.3f (%.3f)  ,  ' % (self.avg_date_time / self.num_frames * batch_size, self.avg_gpu_trans_time / self.num_frames * batch_size)
            print_str += 'ForwardTime: %.3f  ,  ' % (self.avg_forward_time / self.num_frames * batch_size)
            print_str += 'TotalTime: %.3f  ,  ' % ((current_time - self.start_time) / self.num_frames * batch_size)
            # print_str += 'DataTime: %.3f (%.3f)  ,  ' % (self.data_read_done_time - prev_frame_time_backup, self.data_to_gpu_time - self.data_read_done_time)
            # print_str += 'ForwardTime: %.3f  ,  ' % (current_time - self.data_to_gpu_time)
            # print_str += 'TotalTime: %.3f  ,  ' % (current_time - prev_frame_time_backup)

            for name, val in self.stats[loader.name].items():
                if (self.settings.print_stats is None or name in self.settings.print_stats):
                    if name=='Coord': continue
                    if hasattr(val, 'avg'):
                        print_str += '%s: %.5f  ,  ' % (name, val.avg)
                    # else:
                    #     print_str += '%s: %r  ,  ' % (name, val)

            print(print_str[:-5])
            log_str = print_str[:-5] + '\n'
            with open(self.settings.log_file, 'a') as f:
                f.write(log_str)

    def _stats_new_epoch(self):
        # Record learning rate
        for loader in self.loaders:
            if loader.training:
                try:
                    lr_list = self.lr_scheduler.get_last_lr()
                except:
                    lr_list = self.lr_scheduler._get_lr(self.epoch)
                for i, lr in enumerate(lr_list):
                    var_name = 'LearningRate/group{}'.format(i)
                    if var_name not in self.stats[loader.name].keys():
                        self.stats[loader.name][var_name] = StatValue()
                    self.stats[loader.name][var_name].update(lr)

        for loader_stats in self.stats.values():
            if loader_stats is None:
                continue
            for stat_value in loader_stats.values():
                if hasattr(stat_value, 'new_epoch'):
                    stat_value.new_epoch()

    def _write_tensorboard(self):
        if self.epoch == 1:
            self.tensorboard_writer.write_info(self.settings.script_name, self.settings.description)

        self.tensorboard_writer.write_epoch(self.stats, self.epoch)

lib/train/trainers/ltr_trainer.py
TsingWei/LiteTrack
from .tensor import TensorDict, TensorList
from .timeout import timeout, timeout_context, timer

lib/utils/__init__.py
TsingWei/LiteTrack
import torch
from torchvision.ops.boxes import box_area
import numpy as np


def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=-1)

def box_cxcywh_to_xywh(x):
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         w, h]
    return torch.stack(b, dim=-1)


def box_xywh_to_xyxy(x):
    x1, y1, w, h = x.unbind(-1)
    b = [x1, y1, x1 + w, y1 + h]
    return torch.stack(b, dim=-1)

def xywh_to_cxcywh(boxes):
    cx = boxes[:, 0] + boxes[:, 2] / 2
    cy = boxes[:, 1] + boxes[:, 3] / 2
    w = boxes[:, 2]
    h = boxes[:, 3]
    return torch.stack([cx, cy, w, h], dim=1)


def box_xyxy_to_xywh(x):
    x1, y1, x2, y2 = x.unbind(-1)
    b = [x1, y1, x2 - x1, y2 - y1]
    return torch.stack(b, dim=-1)


def box_xyxy_to_cxcywh(x):
    if x is None:
        return None
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2,
         (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


# modified from torchvision to also return the union
'''Note that this function only supports shape (N,4)'''


def box_iou(boxes1, boxes2):
    """

    :param boxes1: (N, 4) (x1,y1,x2,y2)
    :param boxes2: (N, 4) (x1,y1,x2,y2)
    :return:
    """
    area1 = box_area(boxes1) # (N,)
    area2 = box_area(boxes2) # (N,)

    lt = torch.max(boxes1[:, :2], boxes2[:, :2])  # (N,2)
    rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])  # (N,2)

    wh = (rb - lt).clamp(min=0)  # (N,2)
    inter = wh[:, 0] * wh[:, 1]  # (N,)

    union = area1 + area2 - inter

    iou = inter / union
    return iou, union


'''Note that this implementation is different from DETR's'''


def generalized_box_iou(boxes1, boxes2):
    """
    Generalized IoU from https://giou.stanford.edu/

    The boxes should be in [x0, y0, x1, y1] format

    boxes1: (N, 4)
    boxes2: (N, 4)
    """
    # degenerate boxes gives inf / nan results
    # so do an early check
    # try:
    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
    iou, union = box_iou(boxes1, boxes2) # (N,)

    lt = torch.min(boxes1[:, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # (N,2)
    area = wh[:, 0] * wh[:, 1] # (N,)

    return iou - (area - union) / area, iou


def giou_loss(boxes1, boxes2):
    """

    :param boxes1: (N, 4) (x1,y1,x2,y2)
    :param boxes2: (N, 4) (x1,y1,x2,y2)
    :return:
    """
    giou, iou = generalized_box_iou(boxes1, boxes2)
    return (1 - giou), iou


def clip_box(box: list, H, W, margin=0):
    x1, y1, w, h = box
    x2, y2 = x1 + w, y1 + h
    x1 = min(max(0, x1), W-margin)
    x2 = min(max(margin, x2), W)
    y1 = min(max(0, y1), H-margin)
    y2 = min(max(margin, y2), H)
    w = max(margin, x2-x1)
    h = max(margin, y2-y1)
    return [x1, y1, w, h]

# modified from torchvision to also return the union
def box_iou_pairwise(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, :2], boxes2[:, :2])  # [N,2]
    rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])  # [N,2]

    wh = (rb - lt).clamp(min=0)  # [N,2]
    inter = wh[:, 0] * wh[:, 1]  # [N]

    union = area1 + area2 - inter

    iou = inter / union
    return iou, union


def generalized_box_iou_pairwise(boxes1, boxes2):
    """
    Generalized IoU from https://giou.stanford.edu/

    Input:
        - boxes1, boxes2: N,4
    Output:
        - giou: N, 4
    """
    # degenerate boxes gives inf / nan results
    # so do an early check
    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
    assert boxes1.shape == boxes2.shape
    iou, union = box_iou_pairwise(boxes1, boxes2) # N, 4

    lt = torch.min(boxes1[:, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,2]
    area = wh[:, 0] * wh[:, 1]

    return iou - (area - union) / area

def masks_to_boxes(masks):
    """Compute the bounding boxes around the provided masks

    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.

    Returns a [N, 4] tensors, with the boxes in xyxy format
    """
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device)

    h, w = masks.shape[-2:]

    y = torch.arange(0, h, dtype=torch.float)
    x = torch.arange(0, w, dtype=torch.float)
    y, x = torch.meshgrid(y, x)

    x_mask = (masks * x.unsqueeze(0))
    x_max = x_mask.flatten(1).max(-1)[0]
    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    y_mask = (masks * y.unsqueeze(0))
    y_max = y_mask.flatten(1).max(-1)[0]
    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    return torch.stack([x_min, y_min, x_max, y_max], 1)

def jitter_box_in_box(ori_box, num_gen=1):
    box = ori_box.repeat(num_gen,1,1)
    dxy = torch.rand_like(box[:,:,:2]) * box[:,:,2:] - 0.5 * box[:,:,2:]
    new_xy = dxy + box[:,:,:2]
    new_wh = torch.rand_like(box[:,:,2:]) * (box[:,:,2:] - 2*torch.abs(dxy))
    new_box = torch.cat((new_xy,new_wh), -1)
    return new_box

if __name__ == '__main__':
    x = torch.rand(5, 4)
    y = torch.rand(3, 4)
    iou, union = box_iou(x, y)
    import ipdb; ipdb.set_trace()

lib/utils/box_ops.py
TsingWei/LiteTrack
from lib.utils.misc import inverse_sigmoid
import torch, torchvision


def prepare_for_cdn_certain(dn_args, training, num_queries, num_classes, hidden_dim, label_enc):
    """
        The num of group and the num of neg samples within each group are certain.
        1) Box denoising use pooled template feat within the image-pair;
        2) Lable denoising use borrowed pooled feat across the image-pair.
        :param dn_args: targets, dn_number, label_noise_ratio, box_noise_scale
        :param training: if it is training or inference
        :param num_queries: number of queires
        :param num_classes: number of classes
        :param hidden_dim: transformer hidden dim
        :param label_enc: encode labels in dn
        :return:
        """
    if training:
        gt_boxes, dn_number, label_noise_ratio, box_noise_scale = dn_args
        # positive and negative dn queries
        dn_number = dn_number * 2
        # known = [(torch.ones_like(t['labels'])).cuda() for t in targets]
        known = [(torch.ones(num_classes)).long().cuda() for _ in gt_boxes]
        batch_size = len(known)
        known_num = [sum(k) for k in known]
        if int(max(known_num)) == 0:
            dn_number = 1
        else:
            if dn_number >= 100:
                dn_number = dn_number // (int(max(known_num) * 2))
            elif dn_number < 1:
                dn_number = 1
        if dn_number == 0:
            dn_number = 1
        unmask_bbox = unmask_label = torch.cat(known)

        # labels = torch.cat([t['labels'] for t in targets])
        # boxes = torch.cat([t['boxes'] for t in targets])
        # batch_idx = torch.cat([torch.full_like(t['labels'].long(), i) for i, t in enumerate(targets)])
        labels = torch.ones(gt_boxes.shape[0], num_classes).long()
        boxes = gt_boxes
        batch_idx = torch.linspace(start=0, end=len(gt_boxes)-1, steps=len(gt_boxes)).long()

        known_indice = torch.nonzero(unmask_label + unmask_bbox)
        known_indice = known_indice.view(-1)

        known_indice = known_indice.repeat(2 * dn_number, 1).view(-1)
        known_labels = labels.repeat(2 * dn_number, 1).view(-1)
        known_bid = batch_idx.repeat(2 * dn_number, 1).view(-1)
        known_bboxs = boxes.repeat(2 * dn_number, 1)
        known_labels_expaned = known_labels.clone()
        known_bbox_expand = known_bboxs.clone()

        if label_noise_ratio > 0:
            p = torch.rand_like(known_labels_expaned.float())
            chosen_indice = torch.nonzero(p < (label_noise_ratio * 0.5)).view(-1)  # half of bbox prob
            new_label = torch.randint_like(chosen_indice, 0, num_classes)  # randomly put a new one here
            known_labels_expaned.scatter_(0, chosen_indice, new_label)
        single_pad = int(max(known_num))

        pad_size = int(single_pad * 2 * dn_number)
        positive_idx = torch.tensor(range(len(boxes))).long().cuda().unsqueeze(0).repeat(dn_number, 1)
        positive_idx += (torch.tensor(range(dn_number)) * len(boxes) * 2).long().cuda().unsqueeze(1)
        positive_idx = positive_idx.flatten()
        negative_idx = positive_idx + len(boxes)
        if box_noise_scale > 0:
            known_bbox_ = torch.zeros_like(known_bboxs)
            known_bbox_[:, :2] = known_bboxs[:, :2] - known_bboxs[:, 2:] / 2
            known_bbox_[:, 2:] = known_bboxs[:, :2] + known_bboxs[:, 2:] / 2

            diff = torch.zeros_like(known_bboxs)
            diff[:, :2] = known_bboxs[:, 2:] / 2
            diff[:, 2:] = known_bboxs[:, 2:] / 2

            rand_sign = torch.randint_like(known_bboxs, low=0, high=2, dtype=torch.float32) * 2.0 - 1.0
            rand_part = torch.rand_like(known_bboxs)
            rand_part[negative_idx] += 1.0
            rand_part *= rand_sign
            known_bbox_ = known_bbox_ + torch.mul(rand_part,
                                                  diff).cuda() * box_noise_scale
            known_bbox_ = known_bbox_.clamp(min=0.0, max=1.0)
            known_bbox_expand[:, :2] = (known_bbox_[:, :2] + known_bbox_[:, 2:]) / 2
            known_bbox_expand[:, 2:] = known_bbox_[:, 2:] - known_bbox_[:, :2]

        m = known_labels_expaned.long().to('cuda')
        input_label_embed = label_enc(m) # original label to embedding
        #
        input_bbox_embed = inverse_sigmoid(known_bbox_expand)

        padding_label = torch.zeros(pad_size, hidden_dim).cuda()
        padding_bbox = torch.zeros(pad_size, 4).cuda()

        input_query_label = padding_label.repeat(batch_size, 1, 1)
        input_query_bbox = padding_bbox.repeat(batch_size, 1, 1)

        map_known_indice = torch.tensor([]).to('cuda')
        if len(known_num):
            map_known_indice = torch.cat([torch.tensor(range(num)) for num in known_num])  # [1,2, 1,2,3]
            map_known_indice = torch.cat([map_known_indice + single_pad * i for i in range(2 * dn_number)]).long()
        if len(known_bid):
            input_query_label[(known_bid.long(), map_known_indice)] = input_label_embed
            input_query_bbox[(known_bid.long(), map_known_indice)] = input_bbox_embed

        tgt_size = pad_size + num_queries
        attn_mask = torch.ones(tgt_size, tgt_size).to('cuda') < 0
        # match query cannot see the reconstruct
        attn_mask[pad_size:, :pad_size] = True
        # reconstruct cannot see each other
        # torchvision.utils.save_image(attn_mask.long().float(), 'x.png')
        for i in range(dn_number):
            if i == 0:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), single_pad * 2 * (i + 1):pad_size] = True
            if i == dn_number - 1:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), :single_pad * i * 2] = True
            else:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), single_pad * 2 * (i + 1):pad_size] = True
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), :single_pad * 2 * i] = True

        dn_meta = {
            'pad_size': pad_size,
            'num_dn_group': dn_number,
        }
    else:

        input_query_label = None
        input_query_bbox = None
        attn_mask = None
        dn_meta = None

    return input_query_label, input_query_bbox, attn_mask, dn_meta

def prepare_for_cdn(dn_args, training, num_queries, num_classes, hidden_dim, tgt, label_enc=None):
    """
        A major difference of DINO from DN-DETR is that the author process pattern embedding pattern embedding in its detector
        forward function and use learnable tgt embedding, so we change this function a little bit.
        :param dn_args: targets, dn_number, label_noise_ratio, box_noise_scale
        :param training: if it is training or inference
        :param num_queries: number of queires
        :param num_classes: number of classes
        :param hidden_dim: transformer hidden dim
        :param tgt: targets from search, the first is center.
        :param label_enc: encode labels in dn
        :return:
        """
    if training:
        num_tgt = tgt.shape[1]
        gt_boxes, dn_number, label_noise_ratio, box_noise_scale = dn_args
        # positive and negative dn queries
        dn_number = dn_number * 2
        # known = [(torch.ones_like(t['labels'])).cuda() for t in targets]
        known = (torch.ones(len(gt_boxes), num_classes)).long().cuda()
        batch_size = len(known)
        known_num = [1 for k in known]
        if int(max(known_num)) == 0:
            dn_number = 1
        else:
            if dn_number >= 100:
                dn_number = dn_number // (int(max(known_num) * 2))
            elif dn_number < 1:
                dn_number = 1
        if dn_number == 0:
            dn_number = 1
        unmask_bbox = unmask_label = known.flatten(0)

        # labels = torch.cat([t['labels'] for t in targets])
        # boxes = torch.cat([t['boxes'] for t in targets])
        # batch_idx = torch.cat([torch.full_like(t['labels'].long(), i) for i, t in enumerate(targets)])
        labels = torch.ones(gt_boxes.shape[0], num_classes).long()
        boxes = gt_boxes
        batch_idx = torch.linspace(start=0, end=len(gt_boxes)-1, steps=len(gt_boxes)).long()

        known_indice = torch.nonzero(unmask_label + unmask_bbox)
        known_indice = known_indice.view(-1)

        known_indice = known_indice.repeat(2 * dn_number, 1).view(-1)
        known_labels = labels.repeat(2 * dn_number, 1).view(-1)
        known_bid = batch_idx.repeat(2 * dn_number, 1).view(-1)
        known_bboxs = boxes.repeat(2 * dn_number, 1)
        known_labels_expaned = known_labels.clone()
        known_bbox_expand = known_bboxs.clone()

        single_pad = int(max(known_num))
        pad_size = int(single_pad * 2 * dn_number)
        positive_idx = torch.tensor(range(len(boxes))).long().cuda().unsqueeze(0).repeat(dn_number, 1)
        positive_idx += (torch.tensor(range(dn_number)) * len(boxes) * 2).long().cuda().unsqueeze(1)
        positive_idx = positive_idx.flatten()
        negative_idx = positive_idx + len(boxes)

        if label_noise_ratio > 0:
            p = torch.rand_like(known_labels_expaned.float())
            chosen_indice = torch.nonzero(p < (label_noise_ratio * 0.5)).view(-1)  # half of bbox prob
            new_label = torch.randint_like(p, 1, num_tgt)  # randomly put a new one here
            new_label[positive_idx] = 0
            # Simply flip the label to 0.
            known_labels_expaned.scatter_(0, chosen_indice, 0)
        

        
        if box_noise_scale > 0:
            known_bbox_ = torch.zeros_like(known_bboxs)
            known_bbox_[:, :2] = known_bboxs[:, :2] - known_bboxs[:, 2:] / 2
            known_bbox_[:, 2:] = known_bboxs[:, :2] + known_bboxs[:, 2:] / 2

            diff = torch.zeros_like(known_bboxs)
            diff[:, :2] = known_bboxs[:, 2:] / 2
            diff[:, 2:] = known_bboxs[:, 2:] / 2

            rand_sign = torch.randint_like(known_bboxs, low=0, high=2, dtype=torch.float32) * 2.0 - 1.0
            rand_part = torch.rand_like(known_bboxs)
            rand_part[negative_idx] += 2.0
            rand_part *= rand_sign
            known_bbox_ = known_bbox_ + torch.mul(rand_part,
                                                  diff).cuda() * box_noise_scale
            known_bbox_ = known_bbox_.clamp(min=0.0, max=1.0)
            known_bbox_expand[:, :2] = (known_bbox_[:, :2] + known_bbox_[:, 2:]) / 2
            known_bbox_expand[:, 2:] = known_bbox_[:, 2:] - known_bbox_[:, :2]

        m = known_labels_expaned.long().to('cuda')
        input_label_embed = label_enc(m) # original label to embedding
        #
        input_bbox_embed = inverse_sigmoid(known_bbox_expand)

        padding_label = torch.zeros(pad_size, hidden_dim).cuda()
        padding_bbox = torch.zeros(pad_size, 4).cuda()

        input_query_label = padding_label.repeat(batch_size, 1, 1)
        input_query_bbox = padding_bbox.repeat(batch_size, 1, 1)

        map_known_indice = torch.tensor([]).to('cuda')
        if len(known_num):
            map_known_indice = torch.cat([torch.tensor(range(num)) for num in known_num])  # [1,2, 1,2,3]
            map_known_indice = torch.cat([map_known_indice + single_pad * i for i in range(2 * dn_number)]).long()
        if len(known_bid):
            # input_query_label[(known_bid.long(), map_known_indice)] = input_label_embed
            input_query_label[(known_bid.long(), map_known_indice)] = tgt[known_bid.long(),new_label.long()]
            # input_query_label = tgt.repeat(1, pad_size, 1)
            input_query_bbox[(known_bid.long(), map_known_indice)] = input_bbox_embed

        tgt_size = pad_size + num_queries
        attn_mask = torch.ones(tgt_size, tgt_size).to('cuda') < 0
        # match query cannot see the reconstruct
        attn_mask[pad_size:, :pad_size] = True
        # reconstruct cannot see each other
        # torchvision.utils.save_image(attn_mask.long().float(), 'x.png')
        for i in range(dn_number):
            if i == 0:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), single_pad * 2 * (i + 1):pad_size] = True
            if i == dn_number - 1:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), :single_pad * i * 2] = True
            else:
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), single_pad * 2 * (i + 1):pad_size] = True
                attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1), :single_pad * 2 * i] = True

        dn_meta = {
            'pad_size': pad_size,
            'num_dn_group': dn_number,
        }
    else:

        input_query_label = None
        input_query_bbox = None
        attn_mask = None
        dn_meta = None

    return input_query_label, input_query_bbox, attn_mask, dn_meta


def dn_post_process(outputs_class, outputs_coord, dn_meta, aux_loss, _set_aux_loss):
    """
        post process of dn after output from the transformer
        put the dn part in the dn_meta
    """
    if dn_meta and dn_meta['pad_size'] > 0:
        output_known_class = outputs_class[:, :, :dn_meta['pad_size'], :]
        output_known_coord = outputs_coord[:, :, :dn_meta['pad_size'], :]
        outputs_class = outputs_class[:, :, dn_meta['pad_size']:, :]
        outputs_coord = outputs_coord[:, :, dn_meta['pad_size']:, :]
        out = {'pred_logits': output_known_class[-1], 'pred_boxes': output_known_coord[-1]}
        if aux_loss:
            out['aux_outputs'] = _set_aux_loss(output_known_class, output_known_coord)
        dn_meta['output_known_lbs_bboxes'] = out
    return outputs_class, outputs_coord


lib/utils/cdn.py
TsingWei/LiteTrack
import math

import torch
import torch.nn.functional as F


def generate_bbox_mask(bbox_mask, bbox):
    b, h, w = bbox_mask.shape
    for i in range(b):
        bbox_i = bbox[i].cpu().tolist()
        bbox_mask[i, int(bbox_i[1]):int(bbox_i[1] + bbox_i[3] - 1), int(bbox_i[0]):int(bbox_i[0] + bbox_i[2] - 1)] = 1
    return bbox_mask


def generate_mask_cond(cfg, bs, device, gt_bbox, generate_bb_mask_only=False):
    template_size = cfg.DATA.TEMPLATE.SIZE
    stride = cfg.MODEL.BACKBONE.STRIDE
    template_feat_size = template_size // stride

    if generate_bb_mask_only:
        box_mask_z = torch.zeros([bs, template_size, template_size], device=device)
        # box_mask_z_ori = data['template_seg'][0].view(-1, 1, *data['template_seg'].shape[2:])  # (batch, 1, 128, 128)
        box_mask_z = generate_bbox_mask(box_mask_z, gt_bbox * template_size).unsqueeze(1).to(
            torch.float)  # (batch, 1, 128, 128)
        # box_mask_z_vis = box_mask_z.cpu().numpy()
        # box_mask_z = F.interpolate(box_mask_z, scale_factor=1. / cfg.MODEL.BACKBONE.STRIDE, mode='bilinear',
        #                            align_corners=False)
        # box_mask_z = box_mask_z.flatten(1).to(torch.bool)
        return box_mask_z

    if cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'ALL':
        box_mask_z = None
    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_POINT':
        if template_feat_size == 8:
            index = slice(3, 4)
        elif template_feat_size == 11:
            index = slice(5, 6)
        elif template_feat_size == 12:
            index = slice(5, 6)
        elif template_feat_size == 7:
            index = slice(3, 4)
        elif template_feat_size == 14:
            index = slice(6, 7)
        elif template_feat_size == 9:
            index = slice(4, 5)
        elif template_feat_size == 13:
            index = slice(6, 7)
        elif template_feat_size == 16:
            index = slice(7, 8)
        else:
            raise NotImplementedError
        box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
        box_mask_z[:, index, index] = 1
        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_REC':
        # use fixed 4x4 region, 3:5 for 8x8
        # use fixed 4x4 region 5:6 for 12x12
        if template_feat_size == 8:
            index = slice(3, 5)
        elif template_feat_size == 12:
            index = slice(5, 7)
        elif template_feat_size == 7:
            index = slice(3, 4)
        else:
            raise NotImplementedError
        box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
        box_mask_z[:, index, index] = 1
        box_mask_z = box_mask_z.flatten(1).to(torch.bool)

    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'GT_BOX':
        box_mask_z = torch.zeros([bs, template_size, template_size], device=device)
        # box_mask_z_ori = data['template_seg'][0].view(-1, 1, *data['template_seg'].shape[2:])  # (batch, 1, 128, 128)
        box_mask_z = generate_bbox_mask(box_mask_z, gt_bbox * template_size).unsqueeze(1).to(
            torch.float)  # (batch, 1, 128, 128)
        # box_mask_z_vis = box_mask_z.cpu().numpy()
        box_mask_z = F.interpolate(box_mask_z, scale_factor=1. / cfg.MODEL.BACKBONE.STRIDE, mode='bilinear',
                                   align_corners=False)
        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
        # box_mask_z_vis = box_mask_z[:, 0, ...].cpu().numpy()
        # gaussian_maps_vis = generate_heatmap(data['template_anno'], self.cfg.DATA.TEMPLATE.SIZE, self.cfg.MODEL.STRIDE)[0].cpu().numpy()
    else:
        raise NotImplementedError

    return box_mask_z

def generate_mask_cond_search(cfg, bs, device, gt_bbox):
    search_size = cfg.DATA.SEARCH.SIZE
    stride = cfg.MODEL.BACKBONE.STRIDE
    template_feat_size = search_size // stride
    # if cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'ALL':
    #     box_mask_z = None
    # elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_POINT':
    #     if template_feat_size == 8:
    #         index = slice(3, 4)
    #     elif template_feat_size == 11:
    #         index = slice(5, 6)
    #     elif template_feat_size == 12:
    #         index = slice(5, 6)
    #     elif template_feat_size == 7:
    #         index = slice(3, 4)
    #     elif template_feat_size == 14:
    #         index = slice(6, 7)
    #     elif template_feat_size == 9:
    #         index = slice(4, 5)
    #     elif template_feat_size == 13:
    #         index = slice(6, 7)
    #     else:
    #         raise NotImplementedError
    #     box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
    #     box_mask_z[:, index, index] = 1
    #     box_mask_z = box_mask_z.flatten(1).to(torch.bool)
    # elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_REC':
    #     # use fixed 4x4 region, 3:5 for 8x8
    #     # use fixed 4x4 region 5:6 for 12x12
    #     if template_feat_size == 8:
    #         index = slice(3, 5)
    #     elif template_feat_size == 12:
    #         index = slice(5, 7)
    #     elif template_feat_size == 7:
    #         index = slice(3, 4)
    #     else:
    #         raise NotImplementedError
    #     box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
    #     box_mask_z[:, index, index] = 1
    #     box_mask_z = box_mask_z.flatten(1).to(torch.bool)

    if True:
        box_mask_z = torch.zeros([bs, search_size, search_size], device=device)
        # box_mask_z_ori = data['template_seg'][0].view(-1, 1, *data['template_seg'].shape[2:])  # (batch, 1, 128, 128)
        box_mask_z = generate_bbox_mask(box_mask_z, gt_bbox * search_size).unsqueeze(1).to(
            torch.float)  # (batch, 1, 128, 128)
        # box_mask_z_vis = box_mask_z.cpu().numpy()
        box_mask_z = F.interpolate(box_mask_z, scale_factor=1. / cfg.MODEL.BACKBONE.STRIDE, mode='bilinear',
                                   align_corners=False)
        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
        # box_mask_z_vis = box_mask_z[:, 0, ...].cpu().numpy()
        # gaussian_maps_vis = generate_heatmap(data['template_anno'], self.cfg.DATA.TEMPLATE.SIZE, self.cfg.MODEL.STRIDE)[0].cpu().numpy()
    else:
        raise NotImplementedError

    return box_mask_z



def adjust_keep_rate(epoch, warmup_epochs, total_epochs, ITERS_PER_EPOCH, base_keep_rate=0.5, max_keep_rate=1, iters=-1):
    if epoch < warmup_epochs:
        return 1
    if epoch >= total_epochs:
        return base_keep_rate
    if iters == -1:
        iters = epoch * ITERS_PER_EPOCH
    total_iters = ITERS_PER_EPOCH * (total_epochs - warmup_epochs)
    iters = iters - ITERS_PER_EPOCH * warmup_epochs
    keep_rate = base_keep_rate + (max_keep_rate - base_keep_rate) \
        * (math.cos(iters / total_iters * math.pi) + 1) * 0.5

    return keep_rate

lib/utils/ce_utils.py
TsingWei/LiteTrack
from abc import ABC

import torch
import torch.nn as nn
import torch.nn.functional as F


def reduce_loss(loss, reduction):
    """Reduce loss as specified.

    Args:
        loss (Tensor): Elementwise loss tensor.
        reduction (str): Options are "none", "mean" and "sum".

    Return:
        Tensor: Reduced loss tensor.
    """
    reduction_enum = F._Reduction.get_enum(reduction)
    # none: 0, elementwise_mean:1, sum: 2
    if reduction_enum == 0:
        return loss
    elif reduction_enum == 1:
        return loss.mean()
    elif reduction_enum == 2:
        return loss.sum()

def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None):
    """Apply element-wise weight and reduce loss.

    Args:
        loss (Tensor): Element-wise loss.
        weight (Tensor): Element-wise weights.
        reduction (str): Same as built-in losses of PyTorch.
        avg_factor (float): Avarage factor when computing the mean of losses.

    Returns:
        Tensor: Processed loss values.
    """
    # if weight is specified, apply element-wise weight
    if weight is not None:
        loss = loss * weight

    # if avg_factor is not specified, just reduce the loss
    if avg_factor is None:
        loss = reduce_loss(loss, reduction)
    else:
        # if reduction is mean, then average the loss by avg_factor
        if reduction == 'mean':
            loss = loss.sum() / avg_factor
        # if reduction is 'none', then do nothing, otherwise raise an error
        elif reduction != 'none':
            raise ValueError('avg_factor can not be used with reduction="sum"')
    return loss

class FocalLoss(nn.Module, ABC):
    def __init__(self, alpha=2, beta=4):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self, prediction, target):
        positive_index = target.eq(1).float()
        negative_index = target.lt(1).float()

        negative_weights = torch.pow(1 - target, self.beta)
        # clamp min value is set to 1e-12 to maintain the numerical stability
        prediction = torch.clamp(prediction, 1e-7)
        _prediction = torch.clamp(1-prediction, 1e-7)
        positive_loss = torch.log(prediction) * torch.pow(_prediction, self.alpha) * positive_index
        negative_loss = torch.log(_prediction) * torch.pow(prediction,
                                                              self.alpha) * negative_weights * negative_index

        num_positive = positive_index.float().sum()
        positive_loss = positive_loss.sum()
        negative_loss = negative_loss.sum()

        if num_positive == 0:
            loss = -negative_loss
        else:
            loss = -(positive_loss + negative_loss) / num_positive

        return loss

# For corner head
class DistributionFocalLoss(nn.Module, ABC):
    def __init__(self, alpha=2, beta=4):
        super(DistributionFocalLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self, prediction, target):
        positive_index = target.gt(0).float()
        negative_index = target.eq(0).float()
        

        # # negative_weights = 1 for all neg
        # negative_weights = torch.pow(1 - target, self.beta)
        # # clamp min value is set to 1e-12 to maintain the numerical stability
        prediction = torch.clamp(prediction, torch.finfo(prediction.dtype).smallest_normal)
        loss = torch.log(prediction) * target * positive_index
        loss = -loss.sum()
        target_ = torch.clamp(target, torch.finfo(target.dtype).smallest_normal)
        min_loss = torch.log(target_) * target * positive_index
        # _prediction = torch.clamp(1-prediction, 1e-7)
        # positive_loss = torch.log(prediction) * torch.pow(_prediction, self.alpha) * positive_index
        # negative_loss = torch.log(_prediction) * torch.pow(prediction,
        #                                                       self.alpha) * negative_weights * negative_index

        # num_positive = positive_index.float().sum()
        # positive_loss = positive_loss.sum()
        # negative_loss = negative_loss.sum()

        # if num_positive == 0:
        #     loss = -negative_loss
        # else:
        #     loss = -(positive_loss + negative_loss) / num_positive

        return loss + min_loss.sum()    

class DFL(nn.Module, ABC):
    def __init__(self, alpha=2, beta=4):
        super(DFL, self).__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self, 
            pred,
            label,
            weight=None,
            reduction='mean',
            avg_factor=None):
        disl = label.long()
        disr = disl + 1

        wl = disr.float() - label
        wr = label - disl.float()
        # if disl.max()>=16 or disr.max()>=16:
        #     print(' ')
        loss = F.cross_entropy(pred, disl, reduction='none') * wl \
            + F.cross_entropy(pred, disr, reduction='none') * wr
        loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss

class QFL(nn.Module, ABC):
    def __init__(self, alpha=2, beta=4):
        super(QFL, self).__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self,
          pred,          # (n, 1)
          label,         # (n) 0, 1-80: 0 is neg, 1-80 is positive
          score,         # (n) reg target 0-1, only positive is good
          weight=None,
          beta=2.0,
          reduction='mean',
          avg_factor=None):
        # all goes to 0
        pred_sigmoid = pred.sigmoid()
        zerolabel = pred_sigmoid.new_zeros(pred.shape)
        loss = F.binary_cross_entropy_with_logits(
            pred, zerolabel, reduction='none') * pred_sigmoid.pow(beta)

        # label = label - 1
        pos = (label > 0).nonzero().squeeze(1)
        
        # positive goes to bbox quality
        pt = score[pos] - pred_sigmoid[pos, 0]
        loss[pos,0] = F.binary_cross_entropy_with_logits(
            pred[pos,0], score[pos], reduction='none') * pt.pow(beta)

        # loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss

class FL(nn.Module, ABC):
    def __init__(self, alpha=0.25, beta=2):
        super(FL, self).__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self,
          pred,          # (n, 1)
          label,         # (n) 0, 1-80: 0 is neg, 1-80 is positive
          score=None,
          weight=None,
          beta=2.0,
          reduction='mean',
          avg_factor=None):
        # all goes to 0
        pred_sigmoid = pred.sigmoid()
        # zerolabel = pred_sigmoid.new_zeros(pred.shape)
        ce_loss = F.binary_cross_entropy_with_logits(pred, label, weight=score, reduction='none')
        p_t = pred_sigmoid * label + (1 - pred_sigmoid) * (1 - label)
        loss = ce_loss * ((1 - p_t) ** self.beta)
        if self.alpha >= 0:
            alpha_t = self.alpha * label + (1 - self.alpha) * (1 - label)
            loss = alpha_t * loss
        return loss


class old_DistributionFocalLoss(nn.Module, ABC):
    def __init__(self):
        super(DistributionFocalLoss, self).__init__()
    def forward(self, pred, label, size):
        size = int(size)
        bs, n_feat, _ = pred.shape
        pred_map = pred.view(bs, n_feat, size+1, size+1)
        x_distri = pred_map.sum(-2)
        y_distri = pred_map.sum(-1)
        label_grid = size * label
        dis_left = label_grid.long().clamp(max=size-1)
        dis_right = dis_left + 1
        weight_left = dis_right.float() - label_grid
        weight_right = label_grid - dis_left.float()
        # bottom and ceiling
        indexs_X = torch.stack((dis_left.view(bs,2,2)[:, :, 0], dis_right.view(bs,2,2)[:, :, 0]), -1)
        indexs_Y = torch.stack((dis_left.view(bs,2,2)[:, :, 1], dis_right.view(bs,2,2)[:, :, 1]), -1)

        pred_X = x_distri.gather(-1, indexs_X)
        pred_Y = y_distri.gather(-1, indexs_Y)
        target_X = torch.stack((weight_left.view(bs,2,2)[:, :, 0], weight_right.view(bs,2,2)[:, :, 0]), -1)
        target_Y = torch.stack((weight_left.view(bs,2,2)[:, :, 1], weight_right.view(bs,2,2)[:, :, 1]), -1)
        loss = F.cross_entropy(pred_X[:,0], target_X[:,0], reduction='none')  \
            + F.cross_entropy(pred_Y[:,1], target_Y[:,1], reduction='none') \
            + F.cross_entropy(pred_Y[:,0], target_Y[:,0], reduction='none') \
            + F.cross_entropy(pred_Y[:,1], target_Y[:,1], reduction='none') \
            
        return loss.mean()

class LBHinge(nn.Module):
    """Loss that uses a 'hinge' on the lower bound.
    This means that for samples with a label value smaller than the threshold, the loss is zero if the prediction is
    also smaller than that threshold.
    args:
        error_matric:  What base loss to use (MSE by default).
        threshold:  Threshold to use for the hinge.
        clip:  Clip the loss if it is above this value.
    """
    def __init__(self, error_metric=nn.MSELoss(), threshold=None, clip=None):
        super().__init__()
        self.error_metric = error_metric
        self.threshold = threshold if threshold is not None else -100
        self.clip = clip

    def forward(self, prediction, label, target_bb=None):
        negative_mask = (label < self.threshold).float()
        positive_mask = (1.0 - negative_mask)

        prediction = negative_mask * F.relu(prediction) + positive_mask * prediction

        loss = self.error_metric(prediction, positive_mask * label)

        if self.clip is not None:
            loss = torch.min(loss, torch.tensor([self.clip], device=loss.device))
        return loss
lib/utils/focal_loss.py
TsingWei/LiteTrack
import numpy as np
import torch


def get_2d_gaussian_map(boxes, heatmap_size=16, alpha=1.0):
    """
    Generate a 2D gaussian map based on the given box.

    Args:
        box (tensor): 4D tensor of shape (bs, 4) for [cx, cy, w, h].
        S (int): Map size.
        alpha (float): Gaussian distribution factor.

    Returns:
        tensor: 2D gaussian map tensor of shape (S, S).
    """
    bs = boxes.shape[0]
    device = boxes.device
    boxes_sized = boxes*heatmap_size
    cx, cy, w, h = boxes_sized[:, 1], boxes_sized[:, 0], boxes_sized[:, 3], boxes_sized[:, 2]

    # create coordinate grid
    x, y = torch.meshgrid(torch.arange(heatmap_size, dtype=torch.float32, device=device),
                          torch.arange(heatmap_size, dtype=torch.float32, device=device))

    # calculate gaussian distribution
    gauss_map = torch.exp(-((x.unsqueeze(0) - cx.unsqueeze(-1).unsqueeze(-1)) ** 2 / (alpha * w.unsqueeze(-1).unsqueeze(-1)) +
                            (y.unsqueeze(0) - cy.unsqueeze(-1).unsqueeze(-1)) ** 2 / (alpha * h.unsqueeze(-1).unsqueeze(-1))))

    # normalize gauss_map
    # gauss_map_max, _= gauss_map.flatten(1).max(dim=-1)
    # gauss_map_sum = gauss_map.max(dim=(1, 2))
    # gauss_map /= gauss_map_max.view(bs, 1, 1)

    return gauss_map.squeeze()

def generate_cls_map(bboxes, heatmap_size=16):
    # gaussian_maps = []
    # heatmap_size = patch_size // stride
    bs = bboxes.shape[0]
    gt_scoremap = torch.zeros(bs, heatmap_size, heatmap_size).to(bboxes.device)
    bbox = (bboxes * heatmap_size + 0.5).long()
    for i in range(bs):
        gt_scoremap[i, bbox[i,1]:bbox[i,3], bbox[i,0]:bbox[i,2]] = 1
    return gt_scoremap

def generate_heatmap(bboxes, heatmap_size=16):
    """
    Generate ground truth heatmap same as CenterNet
    Args:
        bboxes (torch.Tensor): shape of [num_search, bs, 4]

    Returns:
        gaussian_maps: list of generated heatmap

    """
    gaussian_maps = []
    # heatmap_size = patch_size // stride
    for single_patch_bboxes in bboxes:
        bs = single_patch_bboxes.shape[0]
        gt_scoremap = torch.zeros(bs, heatmap_size, heatmap_size)
        classes = torch.arange(bs).to(torch.long)
        bbox = single_patch_bboxes * heatmap_size
        wh = bbox[:, 2:]
        centers_int = (bbox[:, :2] + wh / 2 - 0.5).round()
        CenterNetHeatMap.generate_score_map(gt_scoremap, classes, wh, centers_int, 0.7)
        gaussian_maps.append(gt_scoremap.to(bbox.device))
    return gaussian_maps

def generate_distribution_heatmap(bboxes, map_size=16):
    """
    Generate ground truth ditribution heatmap
    Args:
        bboxes_xyxy (torch.Tensor): shape of [num_search, bs, 4]

    Returns:
        gaussian_maps: list of generated heatmap

    """
    gaussian_maps = []
    heatmap_size = int(map_size+1)
    for single_patch_bboxes in bboxes:
        bs = single_patch_bboxes.shape[0]
        bbox = single_patch_bboxes * map_size
        tl, dr = bbox[:, :2], bbox[:, 2:]
        wh = tl - dr
        tl_scoremap = torch.zeros(bs, heatmap_size, heatmap_size)
        br_scoremap = torch.zeros(bs, heatmap_size, heatmap_size)
        classes = torch.arange(bs).to(torch.long)
        
        # centers= (bbox[:, :2] + wh / 2).round()
        two_IntegrationHeatMap.generate_score_map(tl_scoremap, classes, wh, tl, 0.6)
        two_IntegrationHeatMap.generate_score_map(br_scoremap, classes, wh, dr, 0.6)
        gaussian_maps.append(torch.stack((tl_scoremap.to(bbox.device), br_scoremap.to(bbox.device)), dim=1))
        # gaussian_maps.append()
    return gaussian_maps

class CenterNetHeatMap(object):
    @staticmethod
    def generate_score_map(fmap, gt_class, gt_wh, centers_int, min_overlap):
        radius = CenterNetHeatMap.get_gaussian_radius(gt_wh, min_overlap)
        radius = torch.clamp_min(radius, 0)
        radius = radius.type(torch.int).cpu().numpy()
        for i in range(gt_class.shape[0]):
            channel_index = gt_class[i]
            CenterNetHeatMap.draw_gaussian(fmap[channel_index], centers_int[i], radius[i])

    @staticmethod
    def get_gaussian_radius(box_size, min_overlap):
        """
        copyed from CornerNet
        box_size (w, h), it could be a torch.Tensor, numpy.ndarray, list or tuple
        notice: we are using a bug-version, please refer to fix bug version in CornerNet
        """
        # box_tensor = torch.Tensor(box_size)
        box_tensor = box_size
        width, height = box_tensor[..., 0], box_tensor[..., 1]

        a1 = 1
        b1 = height + width
        c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
        sq1 = torch.sqrt(b1 ** 2 - 4 * a1 * c1)
        r1 = (b1 + sq1) / 2

        a2 = 4
        b2 = 2 * (height + width)
        c2 = (1 - min_overlap) * width * height
        sq2 = torch.sqrt(b2 ** 2 - 4 * a2 * c2)
        r2 = (b2 + sq2) / 2

        a3 = 4 * min_overlap
        b3 = -2 * min_overlap * (height + width)
        c3 = (min_overlap - 1) * width * height
        sq3 = torch.sqrt(b3 ** 2 - 4 * a3 * c3)
        r3 = (b3 + sq3) / 2

        return torch.min(r1, torch.min(r2, r3))

    @staticmethod
    def gaussian2D(radius, sigma=1):
        # m, n = [(s - 1.) / 2. for s in shape]
        m, n = radius
        y, x = np.ogrid[-m: m + 1, -n: n + 1]

        gauss = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
        gauss[gauss < np.finfo(gauss.dtype).eps * gauss.max()] = 0
        return gauss

    @staticmethod
    def draw_gaussian(fmap, center, radius, k=1):
        diameter = 2 * radius + 1
        gaussian = CenterNetHeatMap.gaussian2D((radius, radius), sigma=diameter / 6)
        gaussian = torch.Tensor(gaussian)
        x, y = int(center[0]), int(center[1])
        height, width = fmap.shape[:2]

        left, right = min(x, radius), min(width - x, radius + 1)
        top, bottom = min(y, radius), min(height - y, radius + 1)

        masked_fmap = fmap[y - top: y + bottom, x - left: x + right]
        masked_gaussian = gaussian[radius - top: radius + bottom, radius - left: radius + right]
        if min(masked_gaussian.shape) > 0 and min(masked_fmap.shape) > 0:
            masked_fmap = torch.max(masked_fmap, masked_gaussian * k)
            fmap[y - top: y + bottom, x - left: x + right] = masked_fmap
        # return fmap

class IntegrationHeatMap(object):
    @staticmethod
    def generate_score_map(fmap, gt_class, gt_wh, centers, min_overlap):
        radius = IntegrationHeatMap.get_gaussian_radius(gt_wh, min_overlap)
        radius = torch.clamp_min(radius, 2) # min to neibghour-like
        radius = radius.type(torch.int).cpu().numpy()
        for i in range(gt_class.shape[0]):
            channel_index = gt_class[i]
            IntegrationHeatMap.draw_gaussian(fmap[channel_index], centers[i], radius[i])

    @staticmethod
    def get_gaussian_radius(box_size, min_overlap):
        """
        copyed from CornerNet
        box_size (w, h), it could be a torch.Tensor, numpy.ndarray, list or tuple
        notice: we are using a bug-version, please refer to fix bug version in CornerNet
        """
        # box_tensor = torch.Tensor(box_size)
        box_tensor = box_size
        width, height = box_tensor[..., 0], box_tensor[..., 1]
        r_X = width * min_overlap
        r_Y = height * min_overlap
        return torch.stack((r_X, r_Y), dim=1)

    @staticmethod
    def gaussian2D(radius, bias=(0,0), sigma=(1,1)):
        # m, n = [(s - 1.) / 2. for s in shape]
        n, m = radius[0], radius[1]
        # y, x = np.ogrid[-m: m + 1, -n: n + 1]
        # x, y = x-bias[0], y-bias[1]
        sx, sy = sigma[0], sigma[1]
        # normal = 1/(2. * np.pi * sx * sy)
        mx, my = bias[0], bias[1]
        x = np.linspace(-m, m)
        y = np.linspace(-n, n)  
        x, y = np.meshgrid(x, y)
        gauss = 1. / (2. * np.pi * sx * sy) * np.exp(-((x - mx)**2. / (2. * sx**2.) + (y - my)**2. / (2. * sy**2.)))
        # gauss = np.exp(-(x * x / (2. * sx**2.) + y * y / (2. * sy**2.)) / (2. * sx * sy)) * normal
        gauss[gauss < np.finfo(gauss.dtype).eps * gauss.max()] = 0
        return gauss

    @staticmethod
    def draw_gaussian(fmap, center, radius, k=1):
        height, width = fmap.shape[:2] #(17,17)
        x, y = center[0], center[1]
        x_bottom, y_bottom = x.long().clamp(max=width-1), y.long()
        # x_ceiling, y_ceiling = x_bottom+1, y_bottom+1
        bias_x = x - x_bottom.float()
        bias_y = y - y_bottom.float()

        diameter = 2 * radius + 1
        gaussian = IntegrationHeatMap.gaussian2D(radius, (bias_x.item(),bias_y.item()), sigma=diameter / 18)
        gaussian = torch.Tensor(gaussian)
        x, y = int(center[0]), int(center[1])
        height, width = fmap.shape[:2]

        left, right = min(x, radius[0]), min(width - x, radius[0] + 1)
        top, bottom = min(y, radius[1]), min(height - y, radius[1] + 1)

        masked_fmap = fmap[y - top: y + bottom, x - left: x + right]
        masked_gaussian = gaussian[radius[1] - top: radius[1] + bottom, radius[0] - left: radius[0] + right]
        if min(masked_gaussian.shape) > 0 and min(masked_fmap.shape) > 0:
            masked_fmap = torch.max(masked_fmap, masked_gaussian * k)
            fmap[y - top: y + bottom, x - left: x + right] = masked_fmap
        # return fmap

            
class two_IntegrationHeatMap(object):
    @staticmethod
    def generate_score_map(fmap, gt_class, wh, centers, overlap):
        # gt_class = batchs
        for i in range(gt_class.shape[0]):
            channel_index = gt_class[i]
            two_IntegrationHeatMap.draw_neighbour(fmap[channel_index], centers[i],)

    @staticmethod
    def draw_neighbour(fmap, center):
        # diameter = 2 * radius + 1
        # gaussian = CenterNetHeatMap.neibhour2D((radius, radius), sigma=diameter / 6)
        # gaussian = torch.Tensor(gaussian)

        height, width = fmap.shape[:2] #(17,17)
        x, y = center[0], center[1]
        x_bottom, y_bottom = x.long().clamp(max=width-2), y.long().clamp(max=width-2)
        x_ceiling, y_ceiling = x_bottom+1, y_bottom+1
        weight_x_bottom = x_ceiling.float() - x
        weight_x_ceiling = x - x_bottom.float()
        weight_y_bottom = y_ceiling.float() - y
        weight_y_ceiling = y - y_bottom.float()


        fmap[y_bottom, x_bottom] = weight_x_bottom * weight_y_bottom
        fmap[y_ceiling, x_bottom] = weight_x_bottom * weight_y_ceiling
        fmap[y_bottom, x_ceiling] = weight_x_ceiling * weight_y_bottom
        fmap[y_ceiling, x_ceiling] = weight_x_ceiling * weight_y_ceiling
        # return fmap



def compute_grids(features, strides):
    """
    grids regret to the input image size
    """
    grids = []
    for level, feature in enumerate(features):
        h, w = feature.size()[-2:]
        shifts_x = torch.arange(
            0, w * strides[level],
            step=strides[level],
            dtype=torch.float32, device=feature.device)
        shifts_y = torch.arange(
            0, h * strides[level],
            step=strides[level],
            dtype=torch.float32, device=feature.device)
        shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
        shift_x = shift_x.reshape(-1)
        shift_y = shift_y.reshape(-1)
        grids_per_level = torch.stack((shift_x, shift_y), dim=1) + \
                          strides[level] // 2
        grids.append(grids_per_level)
    return grids


def get_center3x3(locations, centers, strides, range=3):
    '''
    Inputs:
        locations: M x 2
        centers: N x 2
        strides: M
    '''
    range = (range - 1) / 2
    M, N = locations.shape[0], centers.shape[0]
    locations_expanded = locations.view(M, 1, 2).expand(M, N, 2)  # M x N x 2
    centers_expanded = centers.view(1, N, 2).expand(M, N, 2)  # M x N x 2
    strides_expanded = strides.view(M, 1, 1).expand(M, N, 2)  # M x N
    centers_discret = ((centers_expanded / strides_expanded).int() * strides_expanded).float() + \
                      strides_expanded / 2  # M x N x 2
    dist_x = (locations_expanded[:, :, 0] - centers_discret[:, :, 0]).abs()
    dist_y = (locations_expanded[:, :, 1] - centers_discret[:, :, 1]).abs()
    return (dist_x <= strides_expanded[:, :, 0] * range) & \
           (dist_y <= strides_expanded[:, :, 0] * range)


def get_pred(score_map_ctr, size_map, offset_map, feat_size):
    max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)

    idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
    size = size_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
    offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)

    return size * feat_size, offset


def bbox2distance(points, bbox, max_dis=None, dt=0.1):
    """Decode bounding box based to distances.

    Args:
        points (Tensor): Shape (n, 2), [x, y].
        bbox (Tensor): Shape (n, 4), "xyxy" format
        max_shape (tuple): Shape of the image.

    Returns:
        Tensor: Decoded distances.
    """
    l = points[:, 0] - bbox[:, 0]
    t = points[:, 1] - bbox[:, 1]
    r = bbox[:, 2] - points[:, 0]
    b = bbox[:, 3] - points[:, 1]
    if max_dis is not None:
        l = l.clamp(min=0, max=max_dis - dt)
        t = t.clamp(min=0, max=max_dis - dt)
        r = r.clamp(min=0, max=max_dis - dt)
        b = b.clamp(min=0, max=max_dis - dt)
    return torch.stack([l, t, r, b], -1)

# Not norm to 1
def grid_center_flattened(gt_class_map, feat_size=16):
    '''
    gt_class_map: flattened map (bs, n*n)
    '''
    # feat_size = img_size // stride
    pos_ind = torch.nonzero(gt_class_map)
    idx_y = torch.div(pos_ind[:,1], feat_size , rounding_mode='trunc').to(torch.float) + 0.5
    idx_x = (pos_ind[:,1] % feat_size).to(torch.float) + 0.5
    return torch.stack([idx_x, idx_y], dim=1)

def grid_center_2d(gt_class_map, feat_size=16):
    '''
    gt_class_map: 2d map (bs, n, n)
    '''
    # feat_size = img_size // stride
    pos_ind = torch.nonzero(gt_class_map)
    idx_y = pos_ind[:,1].to(torch.float) + 0.5
    idx_x = pos_ind[:,2].to(torch.float) + 0.5
    return torch.stack([idx_x, idx_y], dim=1)
lib/utils/heapmap_utils.py
TsingWei/LiteTrack
import torch


def merge_template_search(inp_list, return_search=False, return_template=False):
    """NOTICE: search region related features must be in the last place"""
    seq_dict = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0),
                "mask": torch.cat([x["mask"] for x in inp_list], dim=1),
                "pos": torch.cat([x["pos"] for x in inp_list], dim=0)}
    if return_search:
        x = inp_list[-1]
        seq_dict.update({"feat_x": x["feat"], "mask_x": x["mask"], "pos_x": x["pos"]})
    if return_template:
        z = inp_list[0]
        seq_dict.update({"feat_z": z["feat"], "mask_z": z["mask"], "pos_z": z["pos"]})
    return seq_dict

def merge_template_search_DAB(inp_list, return_search=False, return_template=False):
    """NOTICE: search region related features must be in the last place"""
    _feats = [x["feat"][0] for x in inp_list]
    _masks = [x["mask"][0] for x in inp_list]
    _pos = [x["pos"][0] for x in inp_list]
    seq_dict = {"feat": _feats,
                "mask": _masks,
                "pos": _pos}
    if return_search:
        x = inp_list[-1]
        seq_dict.update({"feat_x": x["feat"], "mask_x": x["mask"], "pos_x": x["pos"]})
    if return_template:
        z = inp_list[0]
        seq_dict.update({"feat_z": z["feat"], "mask_z": z["mask"], "pos_z": z["pos"]})
    return seq_dict


def get_qkv(inp_list):
    """The 1st element of the inp_list is about the template,
    the 2nd (the last) element is about the search region"""
    dict_x = inp_list[-1]
    dict_c = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0),
              "mask": torch.cat([x["mask"] for x in inp_list], dim=1),
              "pos": torch.cat([x["pos"] for x in inp_list], dim=0)}  # concatenated dict
    q = dict_x["feat"] + dict_x["pos"]
    k = dict_c["feat"] + dict_c["pos"]
    v = dict_c["feat"]
    key_padding_mask = dict_c["mask"]
    return q, k, v, key_padding_mask

lib/utils/merge.py
TsingWei/LiteTrack
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Misc functions, including distributed helpers.

Mostly copy-paste from torchvision references.
"""
import os
import subprocess
import time
from collections import defaultdict, deque, OrderedDict
import datetime
import pickle
from typing import Optional, List

import torch
import torch.distributed as dist
from torch import Tensor

import numpy as np
import colorsys

# needed due to empty tensor bug in pytorch and torchvision 0.5
import torchvision
vers = torchvision.__version__.split('.')
if int(vers[0]) <= 0 and int(vers[1]) < 7:
    from torchvision.ops import _new_empty_tensor
    from torchvision.ops.misc import _output_size

def align_tensor(inputs, max_len=None):
    if max_len is None:
        max_len = max([len(item) for item in inputs])

    return torch.stack([padding_to(item, max_len) for item in inputs])


def padding_to(inputs, max=300):
    if max is None:
        return inputs
    num_padding = max - len(inputs)
    if inputs.dim() > 1:
        padding = inputs.new_zeros(num_padding,
                                   *inputs.size()[1:],
                                   dtype=inputs.dtype)
    else:
        padding = inputs.new_zeros(num_padding, dtype=inputs.dtype)
    inputs = torch.cat([inputs, padding], dim=0)
    return inputs

class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20, fmt=None):
        if fmt is None:
            fmt = "{median:.4f} ({global_avg:.4f})"
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
        dist.barrier()
        dist.all_reduce(t)
        t = t.tolist()
        self.count = int(t[0])
        self.total = t[1]

    @property
    def median(self):
        d = torch.tensor(list(self.deque))
        return d.median().item()

    @property
    def avg(self):
        d = torch.tensor(list(self.deque), dtype=torch.float32)
        return d.mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deque)

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value)


def all_gather(data):
    """
    Run all_gather on arbitrary picklable data (not necessarily tensors)
    Args:
        data: any picklable object
    Returns:
        list[data]: list of data gathered from each rank
    """
    world_size = get_world_size()
    if world_size == 1:
        return [data]

    # serialized to a Tensor
    buffer = pickle.dumps(data)
    storage = torch.ByteStorage.from_buffer(buffer)
    tensor = torch.ByteTensor(storage).to("cuda")

    # obtain Tensor size of each rank
    local_size = torch.tensor([tensor.numel()], device="cuda")
    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
    dist.all_gather(size_list, local_size)
    size_list = [int(size.item()) for size in size_list]
    max_size = max(size_list)

    # receiving Tensor from all ranks
    # we pad the tensor because torch all_gather does not support
    # gathering tensors of different shapes
    tensor_list = []
    for _ in size_list:
        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
    if local_size != max_size:
        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
        tensor = torch.cat((tensor, padding), dim=0)
    dist.all_gather(tensor_list, tensor)

    data_list = []
    for size, tensor in zip(size_list, tensor_list):
        buffer = tensor.cpu().numpy().tobytes()[:size]
        data_list.append(pickle.loads(buffer))

    return data_list


def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict


class MetricLogger(object):
    def __init__(self, delimiter="\t"):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    def update(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            assert isinstance(v, (float, int))
            self.meters[k].update(v)

    def __getattr__(self, attr):
        if attr in self.meters:
            return self.meters[attr]
        if attr in self.__dict__:
            return self.__dict__[attr]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, attr))

    def __str__(self):
        loss_str = []
        for name, meter in self.meters.items():
            loss_str.append(
                "{}: {}".format(name, str(meter))
            )
        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):
        for meter in self.meters.values():
            meter.synchronize_between_processes()

    def add_meter(self, name, meter):
        self.meters[name] = meter

    def log_every(self, iterable, print_freq, header=None):
        i = 0
        if not header:
            header = ''
        start_time = time.time()
        end = time.time()
        iter_time = SmoothedValue(fmt='{avg:.4f}')
        data_time = SmoothedValue(fmt='{avg:.4f}')
        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
        if torch.cuda.is_available():
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}',
                'max mem: {memory:.0f}'
            ])
        else:
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}'
            ])
        MB = 1024.0 * 1024.0
        for obj in iterable:
            data_time.update(time.time() - end)
            yield obj
            iter_time.update(time.time() - end)
            if i % print_freq == 0 or i == len(iterable) - 1:
                eta_seconds = iter_time.global_avg * (len(iterable) - i)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if torch.cuda.is_available():
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time),
                        memory=torch.cuda.max_memory_allocated() / MB))
                else:
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time)))
            i += 1
            end = time.time()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        print('{} Total time: {} ({:.4f} s / it)'.format(
            header, total_time_str, total_time / len(iterable)))


def get_sha():
    cwd = os.path.dirname(os.path.abspath(__file__))

    def _run(command):
        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
    sha = 'N/A'
    diff = "clean"
    branch = 'N/A'
    try:
        sha = _run(['git', 'rev-parse', 'HEAD'])
        subprocess.check_output(['git', 'diff'], cwd=cwd)
        diff = _run(['git', 'diff-index', 'HEAD'])
        diff = "has uncommited changes" if diff else "clean"
        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
    except Exception:
        pass
    message = f"sha: {sha}, status: {diff}, branch: {branch}"
    return message


def collate_fn(batch):
    batch = list(zip(*batch))
    batch[0] = nested_tensor_from_tensor_list(batch[0])
    return tuple(batch)


def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0] # get the first one
    for sublist in the_list[1:]: # [h,w,3]
        for index, item in enumerate(sublist): # index: 0,1,2
            maxes[index] = max(maxes[index], item) # compare current max with the other elements in the whole
    return maxes


class NestedTensor(object):
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        # type: (Device) -> NestedTensor # noqa
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            assert mask is not None
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):
        return str(self.tensors)


def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    # TODO make this more general
    if tensor_list[0].ndim == 3:
        if torchvision._is_tracing():
            # nested_tensor_from_tensor_list() does not export well to ONNX
            # call _onnx_nested_tensor_from_tensor_list() instead
            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # TODO make it support different-sized images
        max_size = _max_by_axis([list(img.shape) for img in tensor_list]) # [[3,h1,w1], [3,h2,w2], [3,h3,w3], ...]
        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
        batch_shape = [len(tensor_list)] + max_size # ()
        b, c, h, w = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) # copy valid regions of the images to the largest padded base.
            m[: img.shape[1], :img.shape[2]] = False
    else:
        raise ValueError('not supported')
    return NestedTensor(tensor, mask)


# _onnx_nested_tensor_from_tensor_list() is an implementation of
# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
@torch.jit.unused
def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
    max_size = []
    for i in range(tensor_list[0].dim()):
        max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64)
        max_size.append(max_size_i)
    max_size = tuple(max_size)

    # work around for
    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
    # m[: img.shape[1], :img.shape[2]] = False
    # which is not yet supported in onnx
    padded_imgs = []
    padded_masks = []
    for img in tensor_list:
        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
        padded_imgs.append(padded_img)

        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
        padded_masks.append(padded_mask.to(torch.bool))

    tensor = torch.stack(padded_imgs)
    mask = torch.stack(padded_masks)

    return NestedTensor(tensor, mask=mask)


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()


def is_main_process():
    return get_rank() == 0


def save_on_master(*args, **kwargs):
    if is_main_process():
        torch.save(*args, **kwargs)


def init_distributed_mode(args):
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    elif 'SLURM_PROCID' in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    else:
        print('Not using distributed mode')
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)
    args.dist_backend = 'nccl'
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                         world_size=args.world_size, rank=args.rank)
    torch.distributed.barrier()
    setup_for_distributed(args.rank == 0)


@torch.no_grad()
def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    if target.numel() == 0:
        return [torch.zeros([], device=output.device)]
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
    """
    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
    This will eventually be supported natively by PyTorch, and this
    class can go away.
    """
    if float(torchvision.__version__[:3]) < 0.7:
        if input.numel() > 0:
            return torch.nn.functional.interpolate(
                input, size, scale_factor, mode, align_corners
            )

        output_shape = _output_size(2, input, size, scale_factor)
        output_shape = list(input.shape[:-2]) + list(output_shape)
        return _new_empty_tensor(input, output_shape)
    else:
        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)

class color_sys():
    def __init__(self, num_colors) -> None:
        self.num_colors = num_colors
        colors=[]
        for i in np.arange(0., 360., 360. / num_colors):
            hue = i/360.
            lightness = (50 + np.random.rand() * 10)/100.
            saturation = (90 + np.random.rand() * 10)/100.
            colors.append(tuple([int(j*255) for j in colorsys.hls_to_rgb(hue, lightness, saturation)]))
        self.colors = colors

    def __call__(self, idx):
        return self.colors[idx]

def inverse_sigmoid(x, eps=1e-3):
    x = x.clamp(min=0, max=1)
    x1 = x.clamp(min=eps)
    x2 = (1 - x).clamp(min=eps)
    return torch.log(x1/x2)

def clean_state_dict(state_dict):
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        if k[:7] == 'module.':
            k = k[7:]  # remove `module.`
        new_state_dict[k] = v
    return new_state_dict
lib/utils/misc.py
TsingWei/LiteTrack
# ==========================================================
# Modified from mmcv
# ==========================================================
import os, sys
import os.path as osp
import ast
import tempfile
import shutil
from importlib import import_module

from argparse import Action

from addict import Dict
from yapf.yapflib.yapf_api import FormatCode

BASE_KEY = '_base_'
DELETE_KEY = '_delete_'
RESERVED_KEYS = ['filename', 'text', 'pretty_text', 'get', 'dump', 'merge_from_dict']


def check_file_exist(filename, msg_tmpl='file "{}" does not exist'):
    if not osp.isfile(filename):
        raise FileNotFoundError(msg_tmpl.format(filename))

class ConfigDict(Dict):

    def __missing__(self, name):
        raise KeyError(name)

    def __getattr__(self, name):
        try:
            value = super(ConfigDict, self).__getattr__(name)
        except KeyError:
            ex = AttributeError(f"'{self.__class__.__name__}' object has no "
                                f"attribute '{name}'")
        except Exception as e:
            ex = e
        else:
            return value
        raise ex


class SLConfig(object):
    """
    config files.
    only support .py file as config now.

    ref: mmcv.utils.config

    Example:
        >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
        >>> cfg.a
        1
        >>> cfg.b
        {'b1': [0, 1]}
        >>> cfg.b.b1
        [0, 1]
        >>> cfg = Config.fromfile('tests/data/config/a.py')
        >>> cfg.filename
        "/home/kchen/projects/mmcv/tests/data/config/a.py"
        >>> cfg.item4
        'test'
        >>> cfg
        "Config [path: /home/kchen/projects/mmcv/tests/data/config/a.py]: "
        "{'item1': [1, 2], 'item2': {'a': 0}, 'item3': True, 'item4': 'test'}"
    """
    @staticmethod
    def _validate_py_syntax(filename):
        with open(filename) as f:
            content = f.read()
        try:
            ast.parse(content)
        except SyntaxError:
            raise SyntaxError('There are syntax errors in config '
                              f'file {filename}')

    @staticmethod
    def _file2dict(filename):
        filename = osp.abspath(osp.expanduser(filename))
        check_file_exist(filename)
        if filename.lower().endswith('.py'):
            with tempfile.TemporaryDirectory() as temp_config_dir:
                temp_config_file = tempfile.NamedTemporaryFile(
                    dir=temp_config_dir, suffix='.py')
                temp_config_name = osp.basename(temp_config_file.name)
                shutil.copyfile(filename,
                                osp.join(temp_config_dir, temp_config_name))
                temp_module_name = osp.splitext(temp_config_name)[0]
                sys.path.insert(0, temp_config_dir)
                SLConfig._validate_py_syntax(filename)
                mod = import_module(temp_module_name)
                sys.path.pop(0)
                cfg_dict = {
                    name: value
                    for name, value in mod.__dict__.items()
                    if not name.startswith('__')
                }
                # delete imported module
                del sys.modules[temp_module_name]
                # close temp file
                temp_config_file.close()
        elif filename.lower().endswith(('.yml', '.yaml', '.json')):
            from .slio import slload
            cfg_dict = slload(filename)
        else:
            raise IOError('Only py/yml/yaml/json type are supported now!')

        cfg_text = filename + '\n'
        with open(filename, 'r') as f:
            cfg_text += f.read()

        # parse the base file
        if BASE_KEY in cfg_dict:
            cfg_dir = osp.dirname(filename)
            base_filename = cfg_dict.pop(BASE_KEY)
            base_filename = base_filename if isinstance(
                base_filename, list) else [base_filename]

            cfg_dict_list = list()
            cfg_text_list = list()
            for f in base_filename:
                _cfg_dict, _cfg_text = SLConfig._file2dict(osp.join(cfg_dir, f))
                cfg_dict_list.append(_cfg_dict)
                cfg_text_list.append(_cfg_text)

            base_cfg_dict = dict()
            for c in cfg_dict_list:
                if len(base_cfg_dict.keys() & c.keys()) > 0:
                    raise KeyError('Duplicate key is not allowed among bases')
                    # TODO Allow the duplicate key while warnning user
                base_cfg_dict.update(c)

            base_cfg_dict = SLConfig._merge_a_into_b(cfg_dict, base_cfg_dict)
            cfg_dict = base_cfg_dict

            # merge cfg_text
            cfg_text_list.append(cfg_text)
            cfg_text = '\n'.join(cfg_text_list)

        return cfg_dict, cfg_text

    @staticmethod
    def _merge_a_into_b(a, b):
        """merge dict `a` into dict `b` (non-inplace).
            values in `a` will overwrite `b`.
            copy first to avoid inplace modification
            
        Args:
            a ([type]): [description]
            b ([type]): [description]

        Returns:
            [dict]: [description]
        """
        # import ipdb; ipdb.set_trace()
        if not isinstance(a, dict):
            return a

        b = b.copy()
        for k, v in a.items():
            if isinstance(v, dict) and k in b and not v.pop(DELETE_KEY, False):
            
                if not isinstance(b[k], dict) and not isinstance(b[k], list):
                    # if :
                    # import ipdb; ipdb.set_trace()
                    raise TypeError(
                        f'{k}={v} in child config cannot inherit from base '
                        f'because {k} is a dict in the child config but is of '
                        f'type {type(b[k])} in base config. You may set '
                        f'`{DELETE_KEY}=True` to ignore the base config')
                b[k] = SLConfig._merge_a_into_b(v, b[k])
            elif isinstance(b, list):
                try:
                    _ = int(k)
                except:
                    raise TypeError(
                        f'b is a list, '
                        f'index {k} should be an int when input but {type(k)}'
                    )
                b[int(k)] = SLConfig._merge_a_into_b(v, b[int(k)])
            else:   
                b[k] = v
                
        return b

    @staticmethod
    def fromfile(filename):
        cfg_dict, cfg_text = SLConfig._file2dict(filename)
        return SLConfig(cfg_dict, cfg_text=cfg_text, filename=filename)


    def __init__(self, cfg_dict=None, cfg_text=None, filename=None):
        if cfg_dict is None:
            cfg_dict = dict()
        elif not isinstance(cfg_dict, dict):
            raise TypeError('cfg_dict must be a dict, but '
                            f'got {type(cfg_dict)}')
        for key in cfg_dict:
            if key in RESERVED_KEYS:
                raise KeyError(f'{key} is reserved for config file')

        super(SLConfig, self).__setattr__('_cfg_dict', ConfigDict(cfg_dict))
        super(SLConfig, self).__setattr__('_filename', filename)
        if cfg_text:
            text = cfg_text
        elif filename:
            with open(filename, 'r') as f:
                text = f.read()
        else:
            text = ''
        super(SLConfig, self).__setattr__('_text', text)


    @property
    def filename(self):
        return self._filename

    @property
    def text(self):
        return self._text

    @property
    def pretty_text(self):

        indent = 4

        def _indent(s_, num_spaces):
            s = s_.split('\n')
            if len(s) == 1:
                return s_
            first = s.pop(0)
            s = [(num_spaces * ' ') + line for line in s]
            s = '\n'.join(s)
            s = first + '\n' + s
            return s

        def _format_basic_types(k, v, use_mapping=False):
            if isinstance(v, str):
                v_str = f"'{v}'"
            else:
                v_str = str(v)

            if use_mapping:
                k_str = f"'{k}'" if isinstance(k, str) else str(k)
                attr_str = f'{k_str}: {v_str}'
            else:
                attr_str = f'{str(k)}={v_str}'
            attr_str = _indent(attr_str, indent)

            return attr_str

        def _format_list(k, v, use_mapping=False):
            # check if all items in the list are dict
            if all(isinstance(_, dict) for _ in v):
                v_str = '[\n'
                v_str += '\n'.join(
                    f'dict({_indent(_format_dict(v_), indent)}),'
                    for v_ in v).rstrip(',')
                if use_mapping:
                    k_str = f"'{k}'" if isinstance(k, str) else str(k)
                    attr_str = f'{k_str}: {v_str}'
                else:
                    attr_str = f'{str(k)}={v_str}'
                attr_str = _indent(attr_str, indent) + ']'
            else:
                attr_str = _format_basic_types(k, v, use_mapping)
            return attr_str

        def _contain_invalid_identifier(dict_str):
            contain_invalid_identifier = False
            for key_name in dict_str:
                contain_invalid_identifier |= \
                    (not str(key_name).isidentifier())
            return contain_invalid_identifier

        def _format_dict(input_dict, outest_level=False):
            r = ''
            s = []

            use_mapping = _contain_invalid_identifier(input_dict)
            if use_mapping:
                r += '{'
            for idx, (k, v) in enumerate(input_dict.items()):
                is_last = idx >= len(input_dict) - 1
                end = '' if outest_level or is_last else ','
                if isinstance(v, dict):
                    v_str = '\n' + _format_dict(v)
                    if use_mapping:
                        k_str = f"'{k}'" if isinstance(k, str) else str(k)
                        attr_str = f'{k_str}: dict({v_str}'
                    else:
                        attr_str = f'{str(k)}=dict({v_str}'
                    attr_str = _indent(attr_str, indent) + ')' + end
                elif isinstance(v, list):
                    attr_str = _format_list(k, v, use_mapping) + end
                else:
                    attr_str = _format_basic_types(k, v, use_mapping) + end

                s.append(attr_str)
            r += '\n'.join(s)
            if use_mapping:
                r += '}'
            return r

        cfg_dict = self._cfg_dict.to_dict()
        text = _format_dict(cfg_dict, outest_level=True)
        # copied from setup.cfg
        yapf_style = dict(
            based_on_style='pep8',
            blank_line_before_nested_class_or_def=True,
            split_before_expression_after_opening_paren=True)
        text, _ = FormatCode(text, style_config=yapf_style, verify=True)

        return text
    

    def __repr__(self):
        return f'Config (path: {self.filename}): {self._cfg_dict.__repr__()}'

    def __len__(self):
        return len(self._cfg_dict)

    def __getattr__(self, name):
        # # debug
        # print('+'*15)
        # print('name=%s' % name)
        # print("addr:", id(self))
        # # print('type(self):', type(self))
        # print(self.__dict__)
        # print('+'*15)
        # if self.__dict__ == {}:
        #     raise ValueError

        return getattr(self._cfg_dict, name)

    def __getitem__(self, name):
        return self._cfg_dict.__getitem__(name)

    def __setattr__(self, name, value):
        if isinstance(value, dict):
            value = ConfigDict(value)
        self._cfg_dict.__setattr__(name, value)

    def __setitem__(self, name, value):
        if isinstance(value, dict):
            value = ConfigDict(value)
        self._cfg_dict.__setitem__(name, value)

    def __iter__(self):
        return iter(self._cfg_dict)

    def dump(self, file=None):
        # import ipdb; ipdb.set_trace()
        if file is None:
            return self.pretty_text
        else:
            with open(file, 'w') as f:
                f.write(self.pretty_text)

    def merge_from_dict(self, options):
        """Merge list into cfg_dict

        Merge the dict parsed by MultipleKVAction into this cfg.

        Examples:
            >>> options = {'model.backbone.depth': 50,
            ...            'model.backbone.with_cp':True}
            >>> cfg = Config(dict(model=dict(backbone=dict(type='ResNet'))))
            >>> cfg.merge_from_dict(options)
            >>> cfg_dict = super(Config, self).__getattribute__('_cfg_dict')
            >>> assert cfg_dict == dict(
            ...     model=dict(backbone=dict(depth=50, with_cp=True)))

        Args:
            options (dict): dict of configs to merge from.
        """
        option_cfg_dict = {}
        for full_key, v in options.items():
            d = option_cfg_dict
            key_list = full_key.split('.')
            for subkey in key_list[:-1]:
                d.setdefault(subkey, ConfigDict())
                d = d[subkey]
            subkey = key_list[-1]
            d[subkey] = v

        cfg_dict = super(SLConfig, self).__getattribute__('_cfg_dict')
        super(SLConfig, self).__setattr__(
            '_cfg_dict', SLConfig._merge_a_into_b(option_cfg_dict, cfg_dict))

    # for multiprocess
    def __setstate__(self, state):
        self.__init__(state)


    def copy(self):
        return SLConfig(self._cfg_dict.copy())

    def deepcopy(self):
        return SLConfig(self._cfg_dict.deepcopy())


class DictAction(Action):
    """
    argparse action to split an argument into KEY=VALUE form
    on the first = and append to a dictionary. List options should
    be passed as comma separated values, i.e KEY=V1,V2,V3
    """

    @staticmethod
    def _parse_int_float_bool(val):
        try:
            return int(val)
        except ValueError:
            pass
        try:
            return float(val)
        except ValueError:
            pass
        if val.lower() in ['true', 'false']:
            return True if val.lower() == 'true' else False
        if val.lower() in ['none', 'null']:
            return None
        return val

    def __call__(self, parser, namespace, values, option_string=None):
        options = {}
        for kv in values:
            key, val = kv.split('=', maxsplit=1)
            val = [self._parse_int_float_bool(v) for v in val.split(',')]
            if len(val) == 1:
                val = val[0]
            options[key] = val
        setattr(namespace, self.dest, options)


lib/utils/slconfig.py
TsingWei/LiteTrack
import functools
import signal
import contextlib
import time

def timeout(sec):
    """
    timeout decorator
    :param sec: function raise TimeoutError after ? seconds
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapped_func(*args, **kwargs):

            def _handle_timeout(signum, frame):
                err_msg = f'Function {func.__name__} timed out after {sec} seconds'
                raise TimeoutError(err_msg)

            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(sec)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
            return result

        return wrapped_func
    return decorator

@contextlib.contextmanager
def timeout_context(msg=None, timeout_sec=1, log_func=print):
    def _handle_timeout(signum, frame):
        err_msg = f'{msg} timed out after {timeout_sec} seconds'
        raise TimeoutError(err_msg)

    signal.signal(signal.SIGALRM, _handle_timeout)
    signal.alarm(timeout_sec)
    yield
    signal.alarm(0)
    # time_elapsed = time.perf_counter() - begin_time
    # log_func(f"{msg or 'timer'} | {time_elapsed} sec")

@contextlib.contextmanager
def timer_context(msg=None, log_func=print):
    begin_time = time.perf_counter()
    yield
    time_elapsed = time.perf_counter() - begin_time
    log_func(f"{msg or 'timer'} | {time_elapsed} sec")

def timer(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        begin_time = time.perf_counter()
        res = func(*args, **kwargs)
        time_elapsed = time.perf_counter() - begin_time
        print(f'{func.__name__} | {time_elapsed} sec')
        return res
    return wrapper
lib/utils/timeout.py
TsingWei/LiteTrack
import torch
from bytecode import Bytecode, Instr


class get_local(object):
    cache = {}
    is_activate = False

    def __init__(self, varname):
        self.varname = varname

    def __call__(self, func):
        if not type(self).is_activate:
            return func

        type(self).cache[func.__qualname__] = []
        c = Bytecode.from_code(func.__code__)
        extra_code = [
            Instr('STORE_FAST', '_res'),
            Instr('LOAD_FAST', self.varname),
            Instr('STORE_FAST', '_value'),
            Instr('LOAD_FAST', '_res'),
            Instr('LOAD_FAST', '_value'),
            Instr('BUILD_TUPLE', 2),
            Instr('STORE_FAST', '_result_tuple'),
            Instr('LOAD_FAST', '_result_tuple'),
        ]
        c[-1:-1] = extra_code
        func.__code__ = c.to_code()

        def wrapper(*args, **kwargs):
            res, values = func(*args, **kwargs)
            if isinstance(values, torch.Tensor):
                type(self).cache[func.__qualname__].append(values.detach().cpu().numpy())
            elif isinstance(values, list):  # list of Tensor
                type(self).cache[func.__qualname__].append([value.detach().cpu().numpy() for value in values])
            else:
                raise NotImplementedError
            return res

        return wrapper

    @classmethod
    def clear(cls):
        for key in cls.cache.keys():
            cls.cache[key] = []

    @classmethod
    def activate(cls):
        cls.is_activate = True

lib/utils/variable_hook.py
TsingWei/LiteTrack
import torch
import numpy as np


def numpy_to_torch(a: np.ndarray):
    return torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0)
lib/vis/utils.py
TsingWei/LiteTrack
import visdom
import visdom.server
import cv2
import torch
import copy
import numpy as np
from collections import OrderedDict
from enum import Enum

from lib.vis.plotting import overlay_mask, show_image_with_boxes
from lib.vis.utils import numpy_to_torch


class cv_colors(Enum):
    WHITE = (255, 255, 255)[::-1]
    RED = (0, 0, 255)[::-1]
    GREEN = (0, 255, 0)[::-1]
    BLUE = (255, 0, 0)[::-1]
    PURPLE = (247, 44, 200)[::-1]
    ORANGE = (44, 162, 247)[::-1]
    MINT = (239, 255, 66)[::-1]
    YELLOW = (2, 255, 250)[::-1]
    BLACK = (0, 0, 0)[::-1]


def index_to_color(idx):
    return {
        0: cv_colors.GREEN.value,
        1: cv_colors.BLUE.value,
        2: cv_colors.RED.value,
        3: cv_colors.MINT.value,
        4: cv_colors.YELLOW.value,
        5: cv_colors.WHITE.value,
        6: cv_colors.BLACK.value,
    }[idx]


class VisBase:
    def __init__(self, visdom, show_data, title):
        self.visdom = visdom
        self.show_data = show_data
        self.title = title
        self.raw_data = None

    def update(self, data, **kwargs):
        self.save_data(data, **kwargs)

        if self.show_data:
            self.draw_data()

    def save_data(self, data, **kwargs):
        raise NotImplementedError

    def draw_data(self):
        raise NotImplementedError

    def toggle_display(self, new_mode=None):
        if new_mode is not None:
            self.show_data = new_mode
        else:
            self.show_data = not self.show_data

        if self.show_data:
            self.draw_data()
        else:
            self.visdom.close(self.title)


class VisImage(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data):
        data = data.float()
        self.raw_data = data

    def draw_data(self):
        self.visdom.image(self.raw_data.clone(), opts={'title': self.title}, win=self.title)


class VisHeatmap(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data, **kwargs):
        data = data.squeeze().flip(0)
        if kwargs:
            self.raw_data = [data, kwargs]
        else:
            self.raw_data = [data]
        # self.raw_data = data
    def draw_data(self):

        if len(self.raw_data) == 2:
            self.visdom.heatmap(self.raw_data[0].clone(), opts={'title': self.title + ' ' + self.raw_data[1]['caption'], **self.raw_data[1]}, win=self.title)
        else:
            # self.visdom.heatmap(self.raw_data[0].clone(), opts={'title': self.title, 'xmax':1.0, 'xmin':0.0,}, win=self.title)
            self.visdom.heatmap(self.raw_data[0].clone(), opts={'title': self.title}, win=self.title)

        # self.visdom.heatmap(self.raw_data.clone(), opts={'title': self.title}, win=self.title)

class VisHistogram(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data, **kwargs):
        data = data.squeeze()
        if kwargs:
            self.raw_data = [data, kwargs]
        else:
            self.raw_data = [data]
        # self.raw_data = data

    def draw_data(self):
        if len(self.raw_data) == 2:
            self.visdom.bar(self.raw_data[0].clone(), win=self.title,
                            opts={
                                   **self.raw_data[1], 'ymin':0, 'ymax':1})
        else:
            bar = self.visdom.bar(self.raw_data[0].clone(), win=self.title,
                            opts=dict(title=self.title, stacked=False, ytickmin=0.0, ytickmax=1.0)
                            )
            # self.visdom.win_properties(bar, opts=dict(ytickmin=0, ytickmax=1))
        # win = self.visdom.subplots(num_rows=2, num_cols=2)
        # for i in range(4):
        #     row = i // 2
        #     col = i % 2
        #     self.visdom.histogram(self.data[i], win=win, opts=dict(title=f"Dist {i+1}"), row=row, col=col)
        # self.visdom.heatmap(self.raw_data.clone(), opts={'title': self.title}, win=self.title)

# Vis = visdom.Visdom()
# Vis.histogram()

class VisFeaturemap(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)
        self.block_list = None

    def block_list_callback_handler(self, data):
        self.block_list[data['propertyId']]['value'] = data['value']
        self.visdom.properties(self.block_list, opts={'title': 'Featuremap UI'}, win='featuremap_ui')
        self.draw_data()

    def save_data(self, data):
        data = data.view(-1, *data.shape[-2:])
        data = data.flip(1)
        if self.block_list is None:
            self.block_list = []
            self.draw_feat = []
            for i in range(data.shape[0]):
                self.block_list.append({'type': 'checkbox', 'name': 'Channel {:04d}'.format(i), 'value': False})

            self.visdom.properties(self.block_list, opts={'title': 'Featuremap UI'}, win='featuremap_ui')
            self.visdom.register_event_handler(self.block_list_callback_handler, 'featuremap_ui')

        self.raw_data = data

    def draw_data(self):
        if self.block_list is not None and self.show_data:
            for i, d in enumerate(self.block_list):
                if d['value']:
                    fig_title = '{} ch: {:04d}'.format(self.title, i)
                    self.visdom.heatmap(self.raw_data[i, :, :].clone(),
                                        opts={'title': fig_title}, win=fig_title)


class VisCostVolume(VisBase):
    def __init__(self, visdom, show_data, title, flip=False):
        super().__init__(visdom, show_data, title)
        self.show_slice = False
        self.slice_pos = None
        self.flip = flip

    def show_cost_volume(self):
        data = self.raw_data.clone()

        # data_perm = data.permute(2, 0, 3, 1).contiguous()
        data_perm = data.permute(0, 2, 1, 3).contiguous()
        if self.flip:
            data_perm = data_perm.permute(2, 3, 0, 1).contiguous()

        data_perm = data_perm.view(data_perm.shape[0] * data_perm.shape[1], -1)
        self.visdom.heatmap(data_perm.flip(0), opts={'title': self.title}, win=self.title)

    def set_zoom_pos(self, slice_pos):
        self.slice_pos = slice_pos

    def toggle_show_slice(self, new_mode=None):
        if new_mode is not None:
            self.show_slice = new_mode
        else:
            self.show_slice = not self.show_slice

    def show_cost_volume_slice(self):
        slice_pos = self.slice_pos

        # slice_pos: [row, col]
        cost_volume_data = self.raw_data.clone()

        if self.flip:
            cost_volume_slice = cost_volume_data[:, :, slice_pos[0], slice_pos[1]]
        else:
            cost_volume_slice = cost_volume_data[slice_pos[0], slice_pos[1], :, :]
        self.visdom.heatmap(cost_volume_slice.flip(0), opts={'title': self.title}, win=self.title)

    def save_data(self, data):
        data = data.view(data.shape[-2], data.shape[-1], data.shape[-2], data.shape[-1])
        self.raw_data = data

    def draw_data(self):
        if self.show_slice:
            self.show_cost_volume_slice()
        else:
            self.show_cost_volume()


class VisCostVolumeUI(VisBase):
    def cv_ui_handler(self, data):
        zoom_toggled = False
        if data['event_type'] == 'KeyPress':
            if data['key'] == 'ArrowRight':
                self.zoom_pos[1] = min(self.zoom_pos[1] + 1, self.feat_shape[1] - 1)
            elif data['key'] == 'ArrowLeft':
                self.zoom_pos[1] = max(self.zoom_pos[1] - 1, 0)
            elif data['key'] == 'ArrowUp':
                self.zoom_pos[0] = max(self.zoom_pos[0] - 1, 0)
            elif data['key'] == 'ArrowDown':
                self.zoom_pos[0] = min(self.zoom_pos[0] + 1, self.feat_shape[0] - 1)
            elif data['key'] == 'Enter':
                self.zoom_mode = not self.zoom_mode
                zoom_toggled = True

        # Update image
        self.show_image()

        # Update cost volumes
        for block_title, block in self.registered_blocks.items():
            if isinstance(block, VisCostVolume):
                block.set_zoom_pos(self.zoom_pos)
                block.toggle_show_slice(self.zoom_mode)

                if (self.zoom_mode or zoom_toggled) and block.show_data:
                    block.draw_data()

    def __init__(self, visdom, show_data, title, feat_shape, registered_blocks):
        super().__init__(visdom, show_data, title)
        self.feat_shape = feat_shape
        self.zoom_mode = False
        self.zoom_pos = [int((feat_shape[0] - 1) / 2), int((feat_shape[1] - 1) / 2)]
        self.registered_blocks = registered_blocks

        self.visdom.register_event_handler(self.cv_ui_handler, title)

    def draw_grid(self, data):
        stride_r = int(data.shape[1] / self.feat_shape[0])
        stride_c = int(data.shape[2] / self.feat_shape[1])

        # Draw grid
        data[:, list(range(0, data.shape[1], stride_r)), :] = 0
        data[:, :, list(range(0, data.shape[2], stride_c))] = 0

        data[0, list(range(0, data.shape[1], stride_r)), :] = 255
        data[0, :, list(range(0, data.shape[2], stride_c))] = 255

        return data

    def shade_cell(self, data):
        stride_r = int(data.shape[1] / self.feat_shape[0])
        stride_c = int(data.shape[2] / self.feat_shape[1])

        r1 = self.zoom_pos[0] * stride_r
        r2 = min((self.zoom_pos[0] + 1) * stride_r, data.shape[1])

        c1 = self.zoom_pos[1] * stride_c
        c2 = min((self.zoom_pos[1] + 1) * stride_c, data.shape[2])

        factor = 0.8 if self.zoom_mode else 0.5
        data[:, r1:r2, c1:c2] = data[:, r1:r2, c1:c2] * (1 - factor) + torch.tensor([255.0, 0.0, 0.0]).view(3, 1, 1).to(
            data.device) * factor
        return data

    def show_image(self, data=None):
        if data is None:
            data = self.raw_data.clone()

        data = self.draw_grid(data)
        data = self.shade_cell(data)
        self.visdom.image(data, opts={'title': self.title}, win=self.title)

    def save_data(self, data):
        # Ignore feat shape
        data = data[0]
        data = data.float()
        self.raw_data = data

    def draw_data(self):
        self.show_image(self.raw_data.clone())


class VisInfoDict(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)
        self.raw_data = OrderedDict()

    def generate_display_text(self, data):
        display_text = ''
        for key, value in data.items():
            key = key.replace('_', ' ')
            if value is None:
                display_text += '<b>{}</b>: {}<br>'.format(key, 'None')
            elif isinstance(value, (str, int)):
                display_text += '<b>{}</b>: {}<br>'.format(key, value)
            else:
                display_text += '<b>{}</b>: {:.2f}<br>'.format(key, value)

        return display_text

    def save_data(self, data):
        for key, val in data.items():
            self.raw_data[key] = val

    def draw_data(self):
        data = copy.deepcopy(self.raw_data)
        display_text = self.generate_display_text(data)
        self.visdom.text(display_text, opts={'title': self.title}, win=self.title)


class VisText(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data):
        self.raw_data = data

    def draw_data(self):
        data = copy.deepcopy(self.raw_data)
        self.visdom.text(data, opts={'title': self.title}, win=self.title)


class VisLinePlot(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data):
        self.raw_data = data

    def draw_data(self):
        if isinstance(self.raw_data, (list, tuple)):
            data_y = self.raw_data[0].clone()
            data_x = self.raw_data[1].clone()
        else:
            data_y = self.raw_data.clone()
            data_x = torch.arange(data_y.shape[0])

        self.visdom.line(data_y, data_x, opts={'title': self.title}, win=self.title)


class VisTracking(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)
        self.count = 0

    def save_data(self, data, **kwargs):
        image = data[0]
        boxes_masks = data[1:]

        boxes, masks = [], []
        for bm in boxes_masks:
            if bm is None:
                continue
            if isinstance(bm, list):
                boxes.append(torch.Tensor(bm))
                continue
            if len(bm.shape) > 1:
                # Binarize segmentation if a float tensor is provided
                if bm.dtype != np.uint8:
                    bm = (bm > 0.5).astype(np.uint8)
                masks.append(bm)
                continue
            boxes.append(bm.float())

        if kwargs:
            self.raw_data = [image, boxes, masks, kwargs]
        else:
            self.raw_data = [image, boxes, masks]

    def draw_data(self):
        disp_image = self.raw_data[0].copy()

        resize_factor = 1
        if max(disp_image.shape) > 640:
            resize_factor = 640.0 / float(max(disp_image.shape))
            disp_image = cv2.resize(disp_image, None, fx=resize_factor, fy=resize_factor)
            for i, mask in enumerate(self.raw_data[2]):
                self.raw_data[2][i] = cv2.resize(mask, None, fx=resize_factor, fy=resize_factor)

        # if box has score
        scores = None
        if self.raw_data[1][0].shape[0] == 5:
            scores = [box[4].item() if box.shape[0]==5 else 0 for box in self.raw_data[1]]
            self.raw_data[1] = [box[:4] for box in self.raw_data[1]]

        boxes = [resize_factor * b.clone() for b in self.raw_data[1]]

        for i, disp_rect in enumerate(boxes):
            # color = ((255 * ((i % 3) > 0)), 255 * ((i + 1) % 2), (255 * (i % 5)) // 4)
            color = index_to_color(i % 7)
            cv2.rectangle(disp_image,
                          (int(disp_rect[0]), int(disp_rect[1])),
                          (int(disp_rect[0] + disp_rect[2]), int(disp_rect[1] + disp_rect[3])), color, 1)
            if scores is not None and scores[i] != 0:
                cv2.putText(disp_image, "{:.3f}".format(scores[i]), (int(disp_rect[0]), int(disp_rect[1])),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
        for i, mask in enumerate(self.raw_data[2], 1):
            disp_image = overlay_mask(disp_image, mask * i)

        # import os
        # write_img = disp_image.copy()
        # write_img = write_img[:, :, ::-1]
        # # cv2.imwrite(os.path.join('/home/yebotao/test', str(self.count).zfill(3) + '.jpg'), write_img)
        # cv2.imwrite(os.path.join('/home/yebotao/test', self.raw_data[3]['caption'].split('_')[-1] + '.jpg'), write_img)
        # self.count += 1

        disp_image = numpy_to_torch(disp_image).squeeze(0)
        disp_image = disp_image.float()

        if len(self.raw_data) > 3:
            self.visdom.image(disp_image, opts={'title': self.title, **self.raw_data[3]}, win=self.title)
        else:
            self.visdom.image(disp_image, opts={'title': self.title}, win=self.title)
class VisTrackingPts(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)
        self.count = 0

    def save_data(self, data, **kwargs):

        self.raw_data = data

    def draw_data(self):
        disp_image = self.raw_data[0].copy()

        resize_factor = 1


        # if box has score
        scores = self.raw_data[1][4]
        

        box = self.raw_data[1][:4]
        pts = [resize_factor * b for b in self.raw_data[2]]
        for i, disp_rect in enumerate(pts):
            color = index_to_color(1)
            cv2.circle(disp_image, (int(disp_rect[0]),int(disp_rect[1])), 2, color, 2)
        pts = [resize_factor * b for b in self.raw_data[3]]
        for i, disp_rect in enumerate(pts):
            color = index_to_color(3)
            cv2.circle(disp_image, (int(disp_rect[0]),int(disp_rect[1])), 2, color, 2)

        color = index_to_color(0)
        cv2.rectangle(disp_image,
                        (int(box[0]), int(box[1])),
                        (int(box[0] + box[2]), int(box[1] + box[3])), color, 1)
        
        cv2.putText(disp_image, "{:.3f}".format(scores), (int(box[0]), int(box[1])),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

        disp_image = numpy_to_torch(disp_image).squeeze(0)
        disp_image = disp_image.float()

        if False:
            self.visdom.image(disp_image, opts={'title': self.title, **self.raw_data[3]}, win=self.title)
        else:
            self.visdom.image(disp_image, opts={'title': self.title}, win=self.title)


class VisBBReg(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)
        self.block_list = []

    def block_list_callback_handler(self, data):
        self.block_list[data['propertyId']]['value'] = data['value']
        self.visdom.properties(self.block_list, opts={'title': 'BBReg Vis'}, win='bbreg_vis')
        self.draw_data()

    def save_data(self, data):
        self.image = data[0].float()
        self.init_boxes = data[1]
        self.final_boxes = data[2]
        self.final_ious = data[3]

    def draw_data(self):
        if len(self.block_list) == 0:
            self.block_list.append({'type': 'checkbox', 'name': 'ID 0', 'value': True})
            self.block_list.append({'type': 'checkbox', 'name': 'ID 1', 'value': True})
            self.visdom.properties(self.block_list, opts={'title': 'BBReg Vis'}, win='bbreg_vis')
            self.visdom.register_event_handler(self.block_list_callback_handler, 'bbreg_vis')

        disp_image = self.image

        ids = [x['value'] for x in self.block_list]
        init_box_image = show_image_with_boxes(disp_image.clone(), self.init_boxes.clone(), disp_ids=ids)
        final_box_image = show_image_with_boxes(disp_image.clone(), self.final_boxes.clone(), self.final_ious.clone(),
                                                disp_ids=ids)

        self.visdom.image(init_box_image, opts={'title': 'Init Boxes'}, win='Init Boxes')
        self.visdom.image(final_box_image, opts={'title': 'Final Boxes'}, win='Final Boxes')

class VisPoints(VisBase):
    def __init__(self, visdom, show_data, title):
        super().__init__(visdom, show_data, title)

    def save_data(self, data, **kwargs):
        # data = data.squeeze()
        image = data[0]
        pts = data[1:]
        if kwargs:
            self.raw_data = [*data, kwargs]
        else:
            self.raw_data = data
        # self.raw_data = data

    def draw_data(self):
        # a = visdom.Visdom()
        # a.scatter()
        disp_image = self.raw_data[0].copy()

        resize_factor = 1
        if max(disp_image.shape) > 640:
            resize_factor = 640.0 / float(max(disp_image.shape))
            disp_image = cv2.resize(disp_image, None, fx=resize_factor, fy=resize_factor)
            for i, mask in enumerate(self.raw_data[2]):
                self.raw_data[2][i] = cv2.resize(mask, None, fx=resize_factor, fy=resize_factor)

        # if box has score
        # scores = None
        # if self.raw_data[1][0].shape[0] == 5:
        #     scores = [box[4].item() if box.shape[0]==5 else 0 for box in self.raw_data[1]]
        #     self.raw_data[1] = [box[:4] for box in self.raw_data[1]]

        pts = [resize_factor * b for b in self.raw_data[1]]
        # pts = 

        for i, disp_rect in enumerate(pts):
            # color = ((255 * ((i % 3) > 0)), 255 * ((i + 1) % 2), (255 * (i % 5)) // 4)
            color = index_to_color(3)
            cv2.circle(disp_image, (int(disp_rect[0]),int(disp_rect[1])), 4, color, 1)
            # if scores is not None and scores[i] != 0:
            #     cv2.putText(disp_image, "{:.3f}".format(scores[i]), (int(disp_rect[0]), int(disp_rect[1])),
            #                 cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
        # for i, mask in enumerate(self.raw_data[2], 1):
        #     disp_image = overlay_mask(disp_image, mask * i)
        disp_image = numpy_to_torch(disp_image).squeeze(0)
        disp_image = disp_image.float()

        if len(self.raw_data) > 3:
            self.visdom.image(disp_image, opts={'title': self.title, **self.raw_data[3]}, win=self.title)
        else:
            self.visdom.image(disp_image, opts={'title': self.title}, win=self.title)

class Visdom:
    def __init__(self, debug=0, ui_info=None, visdom_info=None, env=None):
        self.debug = debug
        if env is not None:
            self.visdom = visdom.Visdom(server=visdom_info.get('server', '127.0.0.1'),
                                        port=visdom_info.get('port', 8097), env=env)
        else:
            self.visdom = visdom.Visdom(server=visdom_info.get('server', '127.0.0.1'),
                                        port=visdom_info.get('port', 8097))
        self.registered_blocks = {}
        self.blocks_list = []

        self.visdom.properties(self.blocks_list, opts={'title': 'Block List'}, win='block_list')
        self.visdom.register_event_handler(self.block_list_callback_handler, 'block_list')

        # if ui_info is not None:
        #     self.visdom.register_event_handler(ui_info['handler'], ui_info['win_id'])

    def register_event_handler(self, handler, win_id):
         self.visdom.register_event_handler(handler, win_id)

    def block_list_callback_handler(self, data):
        field_name = self.blocks_list[data['propertyId']]['name']

        self.registered_blocks[field_name].toggle_display(data['value'])

        self.blocks_list[data['propertyId']]['value'] = data['value']

        self.visdom.properties(self.blocks_list, opts={'title': 'Block List'}, win='block_list')

    def register(self, data, mode, debug_level=0, title='Data', **kwargs):
        if title not in self.registered_blocks.keys():
            show_data = self.debug >= debug_level

            if title != 'Tracking':
                self.blocks_list.append({'type': 'checkbox', 'name': title, 'value': show_data})

            self.visdom.properties(self.blocks_list, opts={'title': 'Block List'}, win='block_list')

            if mode == 'image':
                self.registered_blocks[title] = VisImage(self.visdom, show_data, title)
            elif mode == 'heatmap':
                self.registered_blocks[title] = VisHeatmap(self.visdom, show_data, title)
            elif mode == 'histogram':
                self.registered_blocks[title] = VisHistogram(self.visdom, show_data, title)
            elif mode == 'cost_volume':
                self.registered_blocks[title] = VisCostVolume(self.visdom, show_data, title)
            elif mode == 'cost_volume_flip':
                self.registered_blocks[title] = VisCostVolume(self.visdom, show_data, title, flip=True)
            elif mode == 'cost_volume_ui':
                self.registered_blocks[title] = VisCostVolumeUI(self.visdom, show_data, title, data[1],
                                                                self.registered_blocks)
            elif mode == 'info_dict':
                self.registered_blocks[title] = VisInfoDict(self.visdom, show_data, title)
            elif mode == 'text':
                self.registered_blocks[title] = VisText(self.visdom, show_data, title)
            elif mode == 'lineplot':
                self.registered_blocks[title] = VisLinePlot(self.visdom, show_data, title)
            elif mode == 'Tracking':
                self.registered_blocks[title] = VisTracking(self.visdom, show_data, title)
            elif mode == 'bbreg':
                self.registered_blocks[title] = VisBBReg(self.visdom, show_data, title)
            elif mode == 'featmap':
                self.registered_blocks[title] = VisFeaturemap(self.visdom, show_data, title)
            elif mode == 'point':
                self.registered_blocks[title] = VisTrackingPts(self.visdom, show_data, title)
            else:
                raise ValueError('Visdom Error: Unknown data mode {}'.format(mode))
        # Update
        self.registered_blocks[title].update(data, **kwargs)

lib/vis/visdom_cus.py
TsingWei/LiteTrack
import _init_paths
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = [8, 8]

from lib.test.analysis.plot_results import plot_results, print_results, print_per_sequence_results
from lib.test.evaluation import get_dataset, trackerlist

trackers = []
dataset_name = 'lasot'
"""stark"""
# trackers.extend(trackerlist(name='stark_s', parameter_name='baseline', dataset_name=dataset_name,
#                             run_ids=None, display_name='STARK-S50'))
# trackers.extend(trackerlist(name='stark_st', parameter_name='baseline', dataset_name=dataset_name,
#                             run_ids=None, display_name='STARK-ST50'))
# trackers.extend(trackerlist(name='stark_st', parameter_name='baseline_R101', dataset_name=dataset_name,
#                             run_ids=None, display_name='STARK-ST101'))
"""TransT"""
# trackers.extend(trackerlist(name='TransT_N2', parameter_name=None, dataset_name=None,
#                             run_ids=None, display_name='TransT_N2', result_only=True))
# trackers.extend(trackerlist(name='TransT_N4', parameter_name=None, dataset_name=None,
#                             run_ids=None, display_name='TransT_N4', result_only=True))
"""pytracking"""
# trackers.extend(trackerlist('atom', 'default', None, range(0,5), 'ATOM'))
# trackers.extend(trackerlist('dimp', 'dimp18', None, range(0,5), 'DiMP18'))
# trackers.extend(trackerlist('dimp', 'dimp50', None, range(0,5), 'DiMP50'))
# trackers.extend(trackerlist('dimp', 'prdimp18', None, range(0,5), 'PrDiMP18'))
# trackers.extend(trackerlist('dimp', 'prdimp50', None, range(0,5), 'PrDiMP50'))
"""ostrack"""
trackers.extend(trackerlist(name='ostrack', parameter_name='cae_center_all_ep300', dataset_name=dataset_name,
                            run_ids=300, display_name='OSTrack256'))
# trackers.extend(trackerlist(name='ostrack', parameter_name='vitb_384_mae_ce_32x4_ep300', dataset_name=dataset_name,
#                             run_ids=None, display_name='OSTrack384'))


dataset = get_dataset(dataset_name)
# dataset = get_dataset('otb', 'nfs', 'uav', 'tc128ce')
# plot_results(trackers, dataset, 'OTB2015', merge_results=True, plot_types=('success', 'norm_prec'),
#              skip_missing_seq=False, force_evaluation=True, plot_bin_gap=0.05)
print_results(trackers, dataset, dataset_name, merge_results=True, plot_types=('success', 'norm_prec', 'prec'))
# print_results(trackers, dataset, 'UNO', merge_results=True, plot_types=('success', 'prec'))

tracking/analysis_results.py
TsingWei/LiteTrack
import os
import sys

prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
    sys.path.append(prj_path)

import argparse
import torch
from thop import profile
from thop.utils import clever_format
import time
import importlib
import copy

def parse_args():
    """
    args for training.
    """
    parser = argparse.ArgumentParser(description='Parse args for training')
    # for train
    parser.add_argument('--script', type=str, default='litetrack', choices=['litetrack'],
                        help='training script name')
    parser.add_argument('--config', type=str, default='B4_cae_center_got10k_ep100', help='yaml configure file name')
    args = parser.parse_args()

    return args


def evaluate_vit(model, template, search):
    '''Speed Test'''
    use_fp16 = False
    if use_fp16:
        template = template.half()
        search = search.half()
        # template_bb = template_bb.half()
    model_ = copy.deepcopy(model)
    device = torch.device(0)
    with torch.autocast(device.type, enabled=use_fp16):
        macs1, params1 = profile(model, inputs=(template, search),
                                custom_ops=None, verbose=False)
    macs, params = clever_format([macs1, params1], "%.3f")
    print('MACs: ', macs)
    print('Params: ', params)
    
    T_w = 200
    T_t = 500
    torch.cuda.synchronize()
    with torch.autocast(device.type, enabled=use_fp16):
        with torch.no_grad():
            # overall
            for i in range(T_w):
                _ = model_(template, search)
            start = time.time()
            for i in range(T_t):
                _ = model_(template, search)
            torch.cuda.synchronize()
            end = time.time()
            avg_lat = (end - start) / T_t
            print("Average latency: %.2f ms" % (avg_lat * 1000))
            print("FPS: %.2f" % (1. / avg_lat))




if __name__ == "__main__":
    device = "cuda:0"
    torch.cuda.set_device(device)
    # Compute the Flops and Params of our STARK-S model
    args = parse_args()
    '''update cfg'''
    yaml_fname = 'experiments/%s/%s.yaml' % (args.script, args.config)
    config_module = importlib.import_module('lib.config.%s.config' % args.script)
    cfg = config_module.cfg
    config_module.update_config_from_file(yaml_fname)
    '''set some values'''
    bs = 1
    z_sz = cfg.DATA.TEMPLATE.SIZE
    x_sz = cfg.DATA.SEARCH.SIZE

    if args.script == "ostrack":
        model_module = importlib.import_module('lib.models')
        model_constructor = model_module.build_ostrack
        model = model_constructor(cfg, training=False)
        # get the template and search
        template = torch.randn(bs, 3, z_sz, z_sz)
        template_bb = torch.tensor([[0.5,0.5,0.5,0.5]])
    
        # template = torch.randn(bs, 64, 768)
        search = torch.randn(bs, 3, x_sz, x_sz)
        # transfer to device
        model = model.to(device)
        model.eval()
        template = template.to(device)
        search = search.to(device)
        template_bb = template_bb.to(device)
        evaluate_vit(model, template, search, template_bb)
 

    if args.script == "litetrack":
        model_module = importlib.import_module('lib.models')
        model_constructor = model_module.build_LiteTrack
        model = model_constructor(cfg, training=False)

        # get the template and search
        
        # template = torch.randn(bs, 3, z_sz, z_sz)
        # template_bb = torch.tensor([[0.5,0.5,0.5,0.5]])
        template = torch.randn(bs, 64, 768)
        search = torch.randn(bs, 3, x_sz, x_sz)

        # transfer to device
        model = model.to(device)
        model.eval()
        template = template.to(device)
        search = search.to(device)

        evaluate_vit(model, template, search)



    else:
        raise NotImplementedError

tracking/profile_model.py
TsingWei/LiteTrack
import os
import sys
import argparse

prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
    sys.path.append(prj_path)

from lib.test.evaluation import get_dataset
from lib.test.evaluation.running import run_dataset
from lib.test.evaluation.tracker import Tracker, trackerlist


def run_tracker(tracker_name, tracker_param, run_id=None, dataset_name='otb', sequence=None, debug=0, threads=0,
                num_gpus=8):
    """Run tracker on sequence or dataset.
    args:
        tracker_name: Name of tracking method.
        tracker_param: Name of parameter file.
        run_id: The run id.
        dataset_name: Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot).
        sequence: Sequence number or name.
        debug: Debug level.
        threads: Number of threads.
    """

    dataset = get_dataset(dataset_name)

    if sequence is not None:
        dataset = [dataset[sequence]]

    # trackers = [Tracker(tracker_name, tracker_param, dataset_name, run_id)]
    trackers = [tuple([tracker_name, tracker_param, dataset_name, ep_id]) for ep_id in run_id]

    run_dataset(dataset, trackers, debug, threads, num_gpus=num_gpus)


def main():
    parser = argparse.ArgumentParser(description='Run tracker on sequence or dataset.')
    parser.add_argument('tracker_name', type=str, help='Name of tracking method.')
    parser.add_argument('tracker_param', type=str, help='Name of config file.')
    parser.add_argument('--runid', type=int, default=None, help='The run id.')
    parser.add_argument('--dataset_name', type=str, default='otb', help='Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot).')
    parser.add_argument('--sequence', type=str, default=None, help='Sequence number or name.')
    parser.add_argument('--debug', type=int, default=0, help='Debug level.')
    parser.add_argument('--threads', type=int, default=0, help='Number of threads.')
    parser.add_argument('--num_gpus', type=int, default=8)
    parser.add_argument('--ep', nargs='+', type=int, default=[100])

    args = parser.parse_args()

    try:
        seq_name = int(args.sequence)
    except:
        seq_name = args.sequence

    run_tracker(args.tracker_name, args.tracker_param, args.ep, args.dataset_name, seq_name, args.debug,
                args.threads, num_gpus=args.num_gpus)


if __name__ == '__main__':
    main()

tracking/test.py
TsingWei/LiteTrack
import os
import sys
import argparse

prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
    sys.path.append(prj_path)

from lib.test.evaluation import get_dataset
from lib.test.evaluation.running import run_dataset
from lib.test.evaluation.tracker import Tracker


def run_tracker(tracker_name, tracker_param, run_id=None, dataset_name='otb', sequence=None, debug=0, threads=0,
                num_gpus=8):
    """Run tracker on sequence or dataset.
    args:
        tracker_name: Name of tracking method.
        tracker_param: Name of parameter file.
        run_id: The run id.
        dataset_name: Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot).
        sequence: Sequence number or name.
        debug: Debug level.
        threads: Number of threads.
    """

    dataset = get_dataset(*dataset_name)

    if sequence is not None:
        dataset = [dataset[sequence]]

    trackers = [Tracker(tracker_name, tracker_param, dataset_name, run_id)]

    run_dataset(dataset, trackers, debug, threads, num_gpus=num_gpus)


def main():
    parser = argparse.ArgumentParser(description='Run tracker on sequence or dataset.')
    parser.add_argument('tracker_name', type=str, help='Name of tracking method.')
    parser.add_argument('tracker_param', type=str, help='Name of config file.')
    parser.add_argument('--runid', type=int, default=None, help='The run id.')
    parser.add_argument('--dataset_name', type=str, default='otb', help='Name of dataset (otb, nfs, uav, tpl, vot, tn, gott, gotv, lasot).')
    parser.add_argument('--sequence', type=str, default=None, help='Sequence number or name.')
    parser.add_argument('--debug', type=int, default=0, help='Debug level.')
    parser.add_argument('--threads', type=int, default=0, help='Number of threads.')
    parser.add_argument('--num_gpus', type=int, default=8)

    args = parser.parse_args()

    try:
        seq_name = int(args.sequence)
    except:
        seq_name = args.sequence

    args.dataset_name = ['trackingnet', 'got10k_test', 'lasot']

    run_tracker(args.tracker_name, args.tracker_param, args.runid, args.dataset_name, seq_name, args.debug,
                args.threads, num_gpus=args.num_gpus)


if __name__ == '__main__':
    main()

tracking/test_exp.py
TsingWei/LiteTrack
import os
import sys

prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
    sys.path.append(prj_path)

import argparse
import torch
from lib.utils.misc import NestedTensor
from thop import profile
from thop.utils import clever_format
import time
import importlib
from thop.vision.basic_hooks import count_linear
import thop
from torch.nn import Linear
import copy
from torch.profiler import profile as t_profile, record_function, ProfilerActivity
# from onnxmltools.utils import float16_converter
import numpy as np
import onnxruntime as ort
from onnxconverter_common import float16
import onnx
# Register the hook for Linear layers


def parse_args():
    """
    args for training.
    """
    parser = argparse.ArgumentParser(description='Parse args for training')
    # for train
    parser.add_argument('--script', type=str, default='ostrack', choices=['ostrack'],
                        help='training script name')
    parser.add_argument('--config', type=str, default='vitb_384_mae_32x4_ep300', help='yaml configure file name')
    args = parser.parse_args()

    return args


def evaluate_vit(model, template, search, template_bb):
    '''Speed Test'''
    model_ = copy.deepcopy(model)
    
    macs1, params1 = profile(model, inputs=(template, search),
                             custom_ops=None, verbose=False)
    macs, params = clever_format([macs1, params1], "%.3f")
    print('overall macs is ', macs)
    print('overall params is ', params)

    output_onnx_name = 'test_net.onnx'

    inputs=(template, search,)
    torch.onnx.export(model, 
        inputs,
        output_onnx_name, 
        input_names=[ "zf", "x"], 
        output_names=["output"],
        opset_version=11,
        export_params=True,
        # verbose=True,
        # dynamic_axes={'input':{0:'batch', 2:'h', 3:'w'}, 'output':{0:'batch', 2:'h2', 3:'w2'}} 
    )

    providers = ['CUDAExecutionProvider']
    model = onnx.load("test_net.onnx")
    dtype = np.float16
    model_fp16 = float16.convert_float_to_float16(model)
    onnx.save(model_fp16, "test_net_fp16.onnx")
    ort_session = ort.InferenceSession("test_net_fp16.onnx", providers=providers)
    output = ort_session.run(output_names=['output'],
                             	input_feed={'zf':  np.array(template.cpu(), dtype=dtype),'x': np.array(search.cpu(), dtype=dtype)}
                                )

    
    T_w = 50
    T_t = 100
    print("testing speed ...")
    # torch.cuda.synchronize()
    # with torch.no_grad():
    #     # overall
    for i in range(T_w):
        output = ort_session.run(output_names=['output'],
                             	input_feed={'zf':  np.array(template.cpu(), dtype=dtype),'x': np.array(search.cpu(), dtype=dtype)}
                                )
    start = time.time()
    for i in range(T_t):
        output = ort_session.run(output_names=['output'],
                             	input_feed={'zf':  np.array(template.cpu(), dtype=dtype),'x': np.array(search.cpu(), dtype=dtype)}
                                )
    torch.cuda.synchronize()
    end = time.time()
    avg_lat = (end - start) / T_t
    print("The average overall latency is %.2f ms" % (avg_lat * 1000))
    print("FPS is %.2f fps" % (1. / avg_lat))
        # for i in range(T_w):
        #     _ = model(template, search)
        # start = time.time()
        # for i in range(T_t):
        #     _ = model(template, search)
        # end = time.time()
        # avg_lat = (end - start) / T_t
        # print("The average backbone latency is %.2f ms" % (avg_lat * 1000))
    # with t_profile(activities=[
    #     ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    #     with record_function("model_inference"):
    #         model_(template, search)
    # print(prof.key_averages(group_by_stack_n=5).table(sort_by="cpu_time_total", row_limit=100))



def get_data(bs, sz):
    img_patch = torch.randn(bs, 3, sz, sz)
    att_mask = torch.rand(bs, sz, sz) > 0.5
    return NestedTensor(img_patch, att_mask)


if __name__ == "__main__":
    device = "cpu"
    # torch.cuda.set_device(device)
    # Compute the Flops and Params of our STARK-S model
    args = parse_args()
    '''update cfg'''
    yaml_fname = 'experiments/%s/%s.yaml' % (args.script, args.config)
    config_module = importlib.import_module('lib.config.%s.config' % args.script)
    cfg = config_module.cfg
    config_module.update_config_from_file(yaml_fname)
    '''set some values'''
    bs = 1
    z_sz = cfg.DATA.TEMPLATE.SIZE
    x_sz = cfg.DATA.SEARCH.SIZE

    if args.script == "ostrack":
        model_module = importlib.import_module('lib.models')
        model_constructor = model_module.build_ostrack
        model = model_constructor(cfg, training=False)
        # get the template and search
        template = torch.randn(bs, 3, z_sz, z_sz)
        template_bb = torch.tensor([[0.5,0.5,0.5,0.5]])
    
        # template = torch.randn(bs, 64, 768)
        search = torch.randn(bs, 3, x_sz, x_sz)
        # transfer to device
        model = model.to(device)
        model.eval()
        template = template.to(device)
        search = search.to(device)
        template_bb = template_bb.to(device)

        merge_layer = cfg.MODEL.BACKBONE.MERGE_LAYER
        if merge_layer <= 0:
            evaluate_vit(model, template, search, template_bb)
        else:
            evaluate_vit_separate(model, template, search)

    else:
        raise NotImplementedError

tracking/torch2onnx.py
TsingWei/LiteTrack
import os
import argparse
import random


def parse_args():
    """
    args for training.
    """
    parser = argparse.ArgumentParser(description='Parse args for training')
    # for train
    parser.add_argument('--script', type=str, help='training script name')
    parser.add_argument('--config', type=str, default='baseline', help='yaml configure file name')
    parser.add_argument('--save_dir', type=str, help='root directory to save checkpoints, logs, and tensorboard')
    parser.add_argument('--mode', type=str, choices=["single", "multiple", "multi_node"], default="multiple",
                        help="train on single gpu or multiple gpus")
    parser.add_argument('--nproc_per_node', type=int, help="number of GPUs per node")  # specify when mode is multiple
    parser.add_argument('--use_lmdb', type=int, choices=[0, 1], default=0)  # whether datasets are in lmdb format
    parser.add_argument('--script_prv', type=str, help='training script name')
    parser.add_argument('--config_prv', type=str, default='baseline', help='yaml configure file name')
    parser.add_argument('--use_wandb', type=int, choices=[0, 1], default=0)  # whether to use wandb
    # for knowledge distillation
    parser.add_argument('--distill', type=int, choices=[0, 1], default=0)  # whether to use knowledge distillation
    parser.add_argument('--script_teacher', type=str, help='teacher script name')
    parser.add_argument('--config_teacher', type=str, help='teacher yaml configure file name')

    # for multiple machines
    parser.add_argument('--rank', type=int, help='Rank of the current process.')
    parser.add_argument('--world-size', type=int, help='Number of processes participating in the job.')
    parser.add_argument('--ip', type=str, default='127.0.0.1', help='IP of the current rank 0.')
    parser.add_argument('--port', type=int, default='20000', help='Port of the current rank 0.')

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    if args.mode == "single":
        train_cmd = "python lib/train/run_training.py --script %s --config %s --save_dir %s --use_lmdb %d " \
                    "--script_prv %s --config_prv %s --distill %d --script_teacher %s --config_teacher %s --use_wandb %d"\
                    % (args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv,
                       args.distill, args.script_teacher, args.config_teacher, args.use_wandb)
    elif args.mode == "multiple":
        train_cmd = "torchrun --standalone --nnodes=1 --nproc_per_node %d --master_port %d lib/train/run_training.py " \
                    "--script %s --config %s --save_dir %s --use_lmdb %d --script_prv %s --config_prv %s --use_wandb %d " \
                    "--distill %d --script_teacher %s --config_teacher %s" \
                    % (args.nproc_per_node, random.randint(10000, 50000), args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv, args.use_wandb,
                       args.distill, args.script_teacher, args.config_teacher)
    elif args.mode == "multi_node":
        train_cmd = "python -m torch.distributed.launch --nproc_per_node %d --master_addr %s --master_port %d --nnodes %d --node_rank %d lib/train/run_training.py " \
                    "--script %s --config %s --save_dir %s --use_lmdb %d --script_prv %s --config_prv %s --use_wandb %d " \
                    "--distill %d --script_teacher %s --config_teacher %s" \
                    % (args.nproc_per_node, args.ip, args.port, args.world_size, args.rank, args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv, args.use_wandb,
                       args.distill, args.script_teacher, args.config_teacher)
    else:
        raise ValueError("mode should be 'single' or 'multiple'.")
    os.system(train_cmd)


if __name__ == "__main__":
    main()

tracking/train.py
TsingWei/LiteTrack
import os
import sys
import argparse

prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
    sys.path.append(prj_path)

from lib.test.evaluation import Tracker


def run_video(tracker_name, tracker_param, videofile, optional_box=None, debug=None, save_results=False):
    """Run the tracker on your webcam.
    args:
        tracker_name: Name of tracking method.
        tracker_param: Name of parameter file.
        debug: Debug level.
    """
    tracker = Tracker(tracker_name, tracker_param, "video",run_id=100)
    tracker.run_video(videofilepath=videofile, optional_box=optional_box, debug=debug, save_results=save_results)


def main():
    parser = argparse.ArgumentParser(description='Run the tracker on your webcam.')
    parser.add_argument('tracker_name', type=str, help='Name of tracking method.')
    parser.add_argument('tracker_param', type=str, help='Name of parameter file.')
    parser.add_argument('videofile', type=str, help='path to a video file.')
    parser.add_argument('--optional_box', type=float, default=None, nargs="+", help='optional_box with format x y w h.')
    parser.add_argument('--debug', type=int, default=0, help='Debug level.')
    parser.add_argument('--save_results', dest='save_results', action='store_true', help='Save bounding boxes')
    parser.set_defaults(save_results=False)

    args = parser.parse_args()

    run_video(args.tracker_name, args.tracker_param, args.videofile, args.optional_box, args.debug, args.save_results)


if __name__ == '__main__':
    main()

tracking/video_demo.py
TsingWei/LiteTrack
import os
import sys
import time

import torch
import numpy as np
import cv2 as cv
from tqdm import tqdm

from lib.vis.visdom_cus import Visdom

env_path = os.path.join(os.path.dirname(__file__), '../lib')
if env_path not in sys.path:
    sys.path.append(env_path)

from lib.test.evaluation import trackerlist, get_dataset
from lib.test.utils.load_text import load_text


class VisResults(object):
    def __init__(self):
        self._init_visdom()

    def vis_dataset(self, dataset, trackers, skip_missing_seq=False, seq_list=[]):
        for seq_id, seq in enumerate(tqdm(dataset)):
            # Load anno
            seq_name = seq.name
            if seq_list:
                if seq_name not in seq_list:
                    continue

            anno_bb = torch.tensor(seq.ground_truth_rect)
            target_visible = torch.tensor(seq.target_visible,
                                          dtype=torch.uint8) if seq.target_visible is not None else None

            all_pred_boxes = []

            for trk_id, trk in enumerate(trackers):
                # Load results
                base_results_path = '{}/{}'.format(trk.results_dir, seq.name)
                results_path = '{}.txt'.format(base_results_path)

                if os.path.isfile(results_path):
                    pred_bb = torch.tensor(load_text(str(results_path), delimiter=('\t', ','), dtype=np.float64))
                    all_pred_boxes.append(pred_bb)
                else:
                    if skip_missing_seq:
                        break
                    else:
                        raise Exception('Result not found. {}'.format(results_path))

            frame_list = seq.frames
            for i in range(len(anno_bb)):
                data = []
                frame = frame_list[i]
                im = cv.imread(frame)
                im = cv.cvtColor(im, cv.COLOR_BGR2RGB)
                # im = torch.from_numpy(im).float().permute(2, 0, 1)
                # im = im.numpy()
                data.append(im)

                gt_box = anno_bb[i]
                data.append(gt_box)
                for tracker_result in all_pred_boxes:
                    data.append(tracker_result[i])

                while self.pause_mode:
                    if self.step:
                        self.step = False
                        break

                if self.next_seq:
                    self.next_seq = False
                    break

                self.update_boxes(data, seq_name + '-' + str(i).zfill(3))
                # self.update_seg_result(im, frame)

    def update_boxes(self, data, caption):
        caption = 'Green: GT, Red: stark_s, Yellow: stark_motion  _' + caption
        self.visdom.register(data, 'Tracking', 1, 'Tracking', caption=caption)

    def update_seg_result(self, frame_img, frame_path):
        seg_mask_path = os.path.join(os.path.dirname(frame_path), 'seg_mask',
                                     os.path.basename(frame_path).replace('jpg', 'png'))
        seg_mask = cv.imread(seg_mask_path)
        alpha = 0.5
        out_img = (alpha * frame_img) + ((1 - alpha) * seg_mask)

        if max(out_img.shape) > 480:
            resize_factor = 480.0 / float(max(out_img.shape))
            out_img = cv.resize(out_img, None, fx=resize_factor, fy=resize_factor)

        out_img = torch.from_numpy(out_img).float().permute(2, 0, 1)
        self.visdom.register(out_img, 'image', 1, 'Segmentation Result')

    def _init_visdom(self, visdom_info=None):
        visdom_info = {} if visdom_info is None else visdom_info
        self.pause_mode = False
        self.step = False
        self.next_seq = False

        try:
            self.visdom = Visdom(1, {'handler': self._visdom_ui_handler, 'win_id': 'Tracking'},
                                 visdom_info=visdom_info, env='vis_results')

            # Show help
            help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \
                        'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \
                        'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \
                        'block list.'
            self.visdom.register(help_text, 'text', 1, 'Help')
        except:
            time.sleep(0.5)
            print('!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n'
                  '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!')

    def _visdom_ui_handler(self, data):
        if data['event_type'] == 'KeyPress':
            if data['key'] == ' ':
                self.pause_mode = not self.pause_mode

            elif data['key'] == 'n':
                self.next_seq = True

            elif data['key'] == 'ArrowRight' and self.pause_mode:
                self.step = True


if __name__ == '__main__':
    viser = VisResults()
    dataset_name = 'lasot'

    trackers = []
    # trackers.extend(trackerlist('defor_stark_s', 'baseline_got10k_only', None, 'defor_stark'))
    # trackers.extend(trackerlist('stark_s', 'baseline_got10k_only', None, 'stark'))
    # trackers.extend(trackerlist('dimp', 'dimp50', dataset_name, None, 'DiMP50'))
    # trackers.extend(trackerlist('sa', 'attn_direct', None, 'SA')) # ori
    # trackers.extend(trackerlist('sa', 'attn1', None, 'SA')) # extrat conv to conver backbone feature from 1024 to 32
    # trackers.extend(trackerlist('sa', 'attn_segaddlabel', None, 'SA')) # add label to seg mask

    # trackers.extend(trackerlist('stark_s', 'baseline_got10k_only_encoder_only_ep500', None, 'STARK-S-Encoder-EP500'))
    # trackers.extend(trackerlist('stark_motion', 'baseline_got10k_only_ep100_002', None, 'STARK-motion-EP100'))
    # trackers.extend(
    #     trackerlist('stark_motion', 'baseline_got10k_only_offset_ep100_001', None, 'STARK-motion-offset-EP100'))

    # trackers.extend(trackerlist('stark_s', 'baseline', dataset_name, None, 'stark_s'))
    # trackers.extend(trackerlist('stark_st', 'baseline', dataset_name, None, 'stark_st'))

    # trackers.extend(trackerlist('stark_mem', 'baseline_roi', dataset_name, None, 'baseline'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi', dataset_name, 17, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi', dataset_name, 32, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi', dataset_name, 36, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ep300', dataset_name, None, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ep300', dataset_name, 6, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ep300', dataset_name, 55, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ep300', dataset_name, 57, 'update_template'))

    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_upsample_syncbn', dataset_name, None, 'baseline_roi_upsample_syncbn'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi', dataset, 9, 'update_template1'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ciou', dataset_name, None, 'update_template'))
    # trackers.extend(trackerlist('stark_mem', 'baseline_roi_ciou', dataset_name, 1, 'update_template'))

    trackers.extend(trackerlist('stark_mem', 'baseline_roi_ep500', dataset_name, None, 'baseline_roi_ep500'))

    trackers.extend(trackerlist('vit_tracker', 'cait_small_224_ep300', dataset_name, None, 'cait_small_224_ep300'))
    trackers.extend(trackerlist('vit_tracker', 'cait_small_224_fcos_new_32x4_ep300', dataset_name, None, 'cait_small_224_fcos_new_32x4_ep300'))
    trackers.extend(trackerlist('vit_tracker', 'cait_small_224_fcos_cn_32x4_ep300_test_cn', dataset_name, None, 'cait_small_224_fcos_cn_32x4_ep300_test_cn'))
    # trackers.extend(trackerlist('vit_tracker', 'cait_small_224_vfloss_64x2_ep100', dataset_name, None, 'cait_small_224_vfloss_64x2_ep100'))
    # trackers.extend(trackerlist('vit_tracker', 'cait_small_224_fcos_cn_32x4_ep300_retest', dataset_name, None, 'retest'))

    dataset = get_dataset(dataset_name)

    # trackers.extend(trackerlist('dimp', 'dimp50', None, 'DiMP50'))
    # trackers.extend(trackerlist('sa', 'attn_direct', None, 'SA'))
    # dataset = get_dataset('vot')
    # 'GOT-10k_Train_001350'
    viser.vis_dataset(dataset, trackers, seq_list=[])
    # viser.vis_dataset(dataset, trackers, seq_list=['GOT-10k_Train_007446'])

tracking/vis_results.py
TsingWei/LiteTrack