First commit.

adet/modeling/vitae_v2/token_transformer.py

+"""
+Borrowed from T2T Transformer
+"""
+import torch
+import torch.nn as nn
+from timm.models.layers import DropPath
+from .NormalCell import Mlp
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, in_dim = None, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., gamma=False, init_values=1e-4):
+        super().__init__()
+        self.num_heads = num_heads
+        self.in_dim = in_dim
+        head_dim = in_dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, in_dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(in_dim, in_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        if gamma:
+            self.gamma1 = nn.Parameter(init_values * torch.ones((in_dim)),requires_grad=True)
+        else:
+            self.gamma1 = 1
+
+    def forward(self, x):
+        B, N, C = x.shape
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.in_dim // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        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, self.in_dim)
+        x = self.proj(x)
+        x = self.proj_drop(self.gamma1 * x)
+        v = v.permute(0, 2, 1, 3).view(B, N, self.in_dim).contiguous()
+        # skip connection
+        x = v + x   # because the original x has different size with current x, use v to do skip connection
+
+        return x
+
+class Token_transformer(nn.Module):
+
+    def __init__(self, dim, in_dim, num_heads, mlp_ratio=1., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, gamma=False, init_values=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, in_dim=in_dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, gamma=gamma, init_values=init_values)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(in_dim)
+        self.mlp = Mlp(in_features=in_dim, hidden_features=int(in_dim*mlp_ratio), out_features=in_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = self.attn(self.norm1(x))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
\ No newline at end of file

adet/modeling/vitae_v2/vitae_v2.py

+from functools import partial
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_
+import numpy as np
+import torch.utils.checkpoint as checkpoint
+from .NormalCell import NormalCell
+from .ReductionCell import ReductionCell
+
+from detectron2.modeling.backbone import Backbone
+from detectron2.modeling.backbone.build import BACKBONE_REGISTRY
+from detectron2.layers import ShapeSpec
+import os.path as osp
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, inter_channel=32, out_channels=48, img_size=None):
+        self.img_size = img_size
+        self.inter_channel = inter_channel
+        self.out_channel = out_channels
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(3, inter_channel, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(inter_channel),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(inter_channel, out_channels, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x, size):
+        x = self.conv3(self.conv2(self.conv1(x)))
+        b, c, h, w = x.shape
+        x = x.permute(0, 2, 3, 1).reshape(b, h * w, c)
+        return x, (h, w)
+
+    def flops(self, ) -> float:
+        flops = 0
+        flops += 3 * self.inter_channel * self.img_size[0] * self.img_size[1] // 4 * 9
+        flops += self.img_size[0] * self.img_size[1] // 4 * self.inter_channel
+        flops += self.inter_channel * self.out_channel * self.img_size[0] * self.img_size[1] // 16 * 9
+        flops += self.img_size[0] * self.img_size[1] // 16 * self.out_channel
+        flops += self.out_channel * self.out_channel * self.img_size[0] * self.img_size[1] // 16
+        return flops
+
+
+class BasicLayer(nn.Module):
+    def __init__(self, img_size=224, in_chans=3, embed_dims=64, token_dims=64, downsample_ratios=4, kernel_size=7, RC_heads=1, NC_heads=6, dilations=[1, 2, 3, 4],
+                RC_op='cat', RC_tokens_type='performer', NC_tokens_type='transformer', RC_group=1, NC_group=64, NC_depth=2, dpr=0.1, mlp_ratio=4., qkv_bias=True,
+                qk_scale=None, drop=0, attn_drop=0., norm_layer=nn.LayerNorm, class_token=False, gamma=False, init_values=1e-4, SE=False, window_size=7,
+                use_checkpoint=False):
+        super().__init__()
+        self.img_size = img_size
+        self.in_chans = in_chans
+        self.embed_dims = embed_dims
+        self.token_dims = token_dims
+        self.downsample_ratios = downsample_ratios
+        self.out_size = self.img_size // self.downsample_ratios
+        self.RC_kernel_size = kernel_size
+        self.RC_heads = RC_heads
+        self.NC_heads = NC_heads
+        self.dilations = dilations
+        self.RC_op = RC_op
+        self.RC_tokens_type = RC_tokens_type
+        self.RC_group = RC_group
+        self.NC_group = NC_group
+        self.NC_depth = NC_depth
+        self.use_checkpoint = use_checkpoint
+        if RC_tokens_type == 'stem':
+            self.RC = PatchEmbedding(inter_channel=token_dims//2, out_channels=token_dims, img_size=img_size)
+        elif downsample_ratios > 1:
+            self.RC = ReductionCell(img_size, in_chans, embed_dims, token_dims, downsample_ratios, kernel_size,
+                            RC_heads, dilations, op=RC_op, tokens_type=RC_tokens_type, group=RC_group, gamma=gamma, init_values=init_values, SE=SE)
+        else:
+            self.RC = nn.Identity()
+        self.NC = nn.ModuleList([
+            NormalCell(token_dims, NC_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop, attn_drop=attn_drop,
+                       drop_path=dpr[i] if isinstance(dpr, list) else dpr, norm_layer=norm_layer, class_token=class_token, group=NC_group, tokens_type=NC_tokens_type,
+                       gamma=gamma, init_values=init_values, SE=SE, img_size=img_size // downsample_ratios, window_size=window_size, shift_size=0)
+        for i in range(NC_depth)])
+
+    def forward(self, x, size):
+        h, w = size
+        x, (h, w) = self.RC(x, (h, w))
+        for nc in self.NC:
+            nc.H = h
+            nc.W = w
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(nc, x)
+            else:
+                x = nc(x)
+        return x, (h, w)
+
+
+class ViTAEv2(Backbone):
+    def __init__(self,
+                img_size=224,
+                in_chans=3,
+                embed_dims=64,
+                token_dims=64,
+                downsample_ratios=[4, 2, 2, 2],
+                kernel_size=[7, 3, 3, 3],
+                RC_heads=[1, 1, 1, 1],
+                NC_heads=4,
+                dilations=[[1, 2, 3, 4], [1, 2, 3], [1, 2], [1, 2]],
+                RC_op='cat',
+                RC_tokens_type='window',
+                NC_tokens_type='transformer',
+                RC_group=[1, 1, 1, 1],
+                NC_group=[1, 32, 64, 64],
+                NC_depth=[2, 2, 6, 2],
+                mlp_ratio=4.,
+                qkv_bias=True,
+                qk_scale=None,
+                drop_rate=0.,
+                attn_drop_rate=0.,
+                drop_path_rate=0.,
+                norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                stages=4,
+                window_size=7,
+                out_indices=(0, 1, 2, 3),
+                frozen_stages=-1,
+                use_checkpoint=False,
+                load_ema=True):
+        super().__init__()
+
+        self.stages = stages
+        self.load_ema = load_ema
+        repeatOrNot = (lambda x, y, z=list: x if isinstance(x, z) else [x for _ in range(y)])
+        self.embed_dims = repeatOrNot(embed_dims, stages)
+        self.tokens_dims = token_dims if isinstance(token_dims, list) else [token_dims * (2 ** i) for i in range(stages)]
+        self.downsample_ratios = repeatOrNot(downsample_ratios, stages)
+        self.kernel_size = repeatOrNot(kernel_size, stages)
+        self.RC_heads = repeatOrNot(RC_heads, stages)
+        self.NC_heads = repeatOrNot(NC_heads, stages)
+        self.dilaions = repeatOrNot(dilations, stages)
+        self.RC_op = repeatOrNot(RC_op, stages)
+        self.RC_tokens_type = repeatOrNot(RC_tokens_type, stages)
+        self.NC_tokens_type = repeatOrNot(NC_tokens_type, stages)
+        self.RC_group = repeatOrNot(RC_group, stages)
+        self.NC_group = repeatOrNot(NC_group, stages)
+        self.NC_depth = repeatOrNot(NC_depth, stages)
+        self.mlp_ratio = repeatOrNot(mlp_ratio, stages)
+        self.qkv_bias = repeatOrNot(qkv_bias, stages)
+        self.qk_scale = repeatOrNot(qk_scale, stages)
+        self.drop = repeatOrNot(drop_rate, stages)
+        self.attn_drop = repeatOrNot(attn_drop_rate, stages)
+        self.norm_layer = repeatOrNot(norm_layer, stages)
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+        self.use_checkpoint = use_checkpoint
+
+        depth = np.sum(self.NC_depth)
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        Layers = []
+        for i in range(stages):
+            startDpr = 0 if i==0 else self.NC_depth[i - 1]
+            Layers.append(
+                BasicLayer(img_size, in_chans, self.embed_dims[i], self.tokens_dims[i], self.downsample_ratios[i],
+                self.kernel_size[i], self.RC_heads[i], self.NC_heads[i], self.dilaions[i], self.RC_op[i],
+                self.RC_tokens_type[i], self.NC_tokens_type[i], self.RC_group[i], self.NC_group[i], self.NC_depth[i], dpr[startDpr:self.NC_depth[i]+startDpr],
+                mlp_ratio=self.mlp_ratio[i], qkv_bias=self.qkv_bias[i], qk_scale=self.qk_scale[i], drop=self.drop[i], attn_drop=self.attn_drop[i],
+                norm_layer=self.norm_layer[i], window_size=window_size, use_checkpoint=use_checkpoint)
+            )
+            img_size = img_size // self.downsample_ratios[i]
+            in_chans = self.tokens_dims[i]
+
+        self.layers = nn.ModuleList(Layers)
+        self.num_layers = len(Layers)
+
+        self._freeze_stages()
+
+        self._out_features = ["stage3", "stage4", "stage5"]
+        self.init_weights()
+
+    def _freeze_stages(self):
+
+        if self.frozen_stages > 0:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+    def init_weights(self):
+        """Initialize the weights in backbone.
+
+        Args:
+            pretrained (str, optional): Path to pre-trained weights.
+                Defaults to None.
+        """
+
+        def _init_weights(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) or isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.bias, 0)
+                nn.init.constant_(m.weight, 1.0)
+
+        self.apply(_init_weights)
+
+    def forward(self, x):
+        """Forward function."""
+        outs = {}
+        b, _, h, w = x.shape
+        for idx, layer in enumerate(self.layers):
+            x, (h, w) = layer(x, (h, w))
+
+            stage_name = "stage" + str(idx + 2)
+            if stage_name in self._out_features:
+                outs[stage_name] = x.reshape(b, h, w, -1).permute(0, 3, 1, 2).contiguous()
+        return outs
+
+    def output_shape(self):
+        return {
+            "stage3": ShapeSpec(channels=128, stride=8),
+            "stage4": ShapeSpec(channels=256, stride=16),
+            "stage5": ShapeSpec(channels=512, stride=32),
+        }
+
+
+@BACKBONE_REGISTRY.register()
+def build_vitaev2_backbone(cfg, input_shape):
+    vitaev2_type = cfg.MODEL.ViTAEv2.TYPE
+
+    if vitaev2_type == 'vitaev2_s':
+        backbone = ViTAEv2(
+            in_chans=3,
+            RC_tokens_type=['window', 'window', 'transformer', 'transformer'],
+            NC_tokens_type=['window', 'window', 'transformer', 'transformer'],
+            embed_dims=[64, 64, 128, 256],
+            token_dims=[64, 128, 256, 512],
+            downsample_ratios=[4, 2, 2, 2],
+            NC_depth=[2, 2, 8, 2],
+            NC_heads=[1, 2, 4, 8],
+            RC_heads=[1, 1, 2, 4],
+            mlp_ratio=4.,
+            NC_group=[1, 32, 64, 128],
+            RC_group=[1, 16, 32, 64],
+            use_checkpoint=True,
+            drop_rate=0.,
+            attn_drop_rate=0.,
+            window_size=7,
+            drop_path_rate=cfg.MODEL.ViTAEv2.DROP_PATH_RATE,
+        )
+    else:
+        raise NotImplementedError
+
+    return backbone
\ No newline at end of file

adet/modeling/vitae_v2/window.py

+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu
+# --------------------------------------------------------
+# Modified by Qiming Zhang
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+
+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
+
+
+def window_partition(x, window_size):
+    """
+    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, H, W):
+    """
+    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 WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, in_dim, out_dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.in_dim = in_dim
+        self.dim = out_dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = out_dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        # self.relative_position_bias_table = nn.Parameter(
+        #     torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # # get pair-wise relative position index for each token inside the window
+        # coords_h = torch.arange(self.window_size[0])
+        # coords_w = torch.arange(self.window_size[1])
+        # coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        # coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        # relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        # relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        # relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        # relative_coords[:, :, 1] += self.window_size[1] - 1
+        # relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        # relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        # self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(in_dim, out_dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(out_dim, out_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        # trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        # relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+        #     self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        # relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        # attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        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
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class WindowTransformerBlock(nn.Module):
+    r""" Window Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, in_dim, out_dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.in_dim = in_dim
+        self.dim = out_dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(in_dim)
+        self.attn = WindowAttention(
+            in_dim=in_dim, out_dim=out_dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(out_dim)
+        mlp_hidden_dim = int(out_dim * mlp_ratio)
+        self.mlp = Mlp(in_features=out_dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
\ No newline at end of file

adet/utils/comm.py

+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from detectron2.utils.comm import get_world_size
+
+
+def reduce_sum(tensor):
+    world_size = get_world_size()
+    if world_size < 2:
+        return tensor
+    tensor = tensor.clone()
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    return tensor
+
+
+def reduce_mean(tensor):
+    num_gpus = get_world_size()
+    total = reduce_sum(tensor)
+    return total.float() / num_gpus
+
+
+def aligned_bilinear(tensor, factor):
+    assert tensor.dim() == 4
+    assert factor >= 1
+    assert int(factor) == factor
+
+    if factor == 1:
+        return tensor
+
+    h, w = tensor.size()[2:]
+    tensor = F.pad(tensor, pad=(0, 1, 0, 1), mode="replicate")
+    oh = factor * h + 1
+    ow = factor * w + 1
+    tensor = F.interpolate(
+        tensor, size=(oh, ow),
+        mode='bilinear',
+        align_corners=True
+    )
+    tensor = F.pad(
+        tensor, pad=(factor // 2, 0, factor // 2, 0),
+        mode="replicate"
+    )
+
+    return tensor[:, :, :oh - 1, :ow - 1]
+
+
+def compute_locations(h, w, stride, device):
+    shifts_x = torch.arange(
+        0, w * stride, step=stride,
+        dtype=torch.float32, device=device
+    )
+    shifts_y = torch.arange(
+        0, h * stride, step=stride,
+        dtype=torch.float32, device=device
+    )
+    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
+    shift_x = shift_x.reshape(-1)
+    shift_y = shift_y.reshape(-1)
+    locations = torch.stack((shift_x, shift_y), dim=1) + stride // 2
+    return locations
+
+
+def compute_ious(pred, target):
+    """
+    Args:
+        pred: Nx4 predicted bounding boxes
+        target: Nx4 target bounding boxes
+        Both are in the form of FCOS prediction (l, t, r, b)
+    """
+    pred_left = pred[:, 0]
+    pred_top = pred[:, 1]
+    pred_right = pred[:, 2]
+    pred_bottom = pred[:, 3]
+
+    target_left = target[:, 0]
+    target_top = target[:, 1]
+    target_right = target[:, 2]
+    target_bottom = target[:, 3]
+
+    target_aera = (target_left + target_right) * \
+                  (target_top + target_bottom)
+    pred_aera = (pred_left + pred_right) * \
+                (pred_top + pred_bottom)
+
+    w_intersect = torch.min(pred_left, target_left) + \
+                  torch.min(pred_right, target_right)
+    h_intersect = torch.min(pred_bottom, target_bottom) + \
+                  torch.min(pred_top, target_top)
+
+    g_w_intersect = torch.max(pred_left, target_left) + \
+                    torch.max(pred_right, target_right)
+    g_h_intersect = torch.max(pred_bottom, target_bottom) + \
+                    torch.max(pred_top, target_top)
+    ac_uion = g_w_intersect * g_h_intersect
+
+    area_intersect = w_intersect * h_intersect
+    area_union = target_aera + pred_aera - area_intersect
+
+    ious = (area_intersect + 1.0) / (area_union + 1.0)
+    gious = ious - (ac_uion - area_union) / ac_uion
+
+    return ious, gious

adet/utils/curve_utils.py

+# borrow from https://github.com/voldemortX/pytorch-auto-drive/blob/master/utils/curve_utils.py
+
+import torch
+import numpy as np
+from scipy.interpolate import splprep, splev
+from scipy.special import comb as n_over_k
+
+
+def upcast(t):
+    # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
+    # https://github.com/pytorch/vision/pull/3383
+    if t.is_floating_point():
+        return t if t.dtype in (torch.float32, torch.float64) else t.float()
+    else:
+        return t if t.dtype in (torch.int32, torch.int64) else t.int()
+
+
+class BezierCurve(object):
+    # Define Bezier curves for curve fitting
+    def __init__(self, order, num_sample_points=50):
+        self.num_point = order + 1
+        self.control_points = []
+        self.bezier_coeff = self.get_bezier_coefficient()
+        self.num_sample_points = num_sample_points
+        self.c_matrix = self.get_bernstein_matrix()
+
+    def get_bezier_coefficient(self):
+        Mtk = lambda n, t, k: t ** k * (1 - t) ** (n - k) * n_over_k(n, k)
+        BezierCoeff = lambda ts: [[Mtk(self.num_point - 1, t, k) for k in range(self.num_point)] for t in ts]
+
+        return BezierCoeff
+
+    def interpolate_lane(self, x, y, n=50):
+        # Spline interpolation of a lane. Used on the predictions
+        assert len(x) == len(y)
+
+        tck, _ = splprep([x, y], s=0, t=n, k=min(3, len(x) - 1))
+
+        u = np.linspace(0., 1., n)
+        return np.array(splev(u, tck)).T
+
+    def get_control_points(self, x, y, interpolate=False):
+        if interpolate:
+            points = self.interpolate_lane(x, y)
+            x = np.array([x for x, _ in points])
+            y = np.array([y for _, y in points])
+
+        middle_points = self.get_middle_control_points(x, y)
+        for idx in range(0, len(middle_points) - 1, 2):
+            self.control_points.append([middle_points[idx], middle_points[idx + 1]])
+
+    def get_bernstein_matrix(self):
+        tokens = np.linspace(0, 1, self.num_sample_points)
+        c_matrix = self.bezier_coeff(tokens)
+        return np.array(c_matrix)
+
+    def save_control_points(self):
+        return self.control_points
+
+    def assign_control_points(self, control_points):
+        self.control_points = control_points
+
+    def quick_sample_point(self, image_size=None):
+        control_points_matrix = np.array(self.control_points)
+        sample_points = self.c_matrix.dot(control_points_matrix)
+        if image_size is not None:
+            sample_points[:, 0] = sample_points[:, 0] * image_size[-1]
+            sample_points[:, -1] = sample_points[:, -1] * image_size[0]
+        return sample_points
+
+    def get_sample_point(self, n=50, image_size=None):
+        '''
+            :param n: the number of sampled points
+            :return: a list of sampled points
+        '''
+        t = np.linspace(0, 1, n)
+        coeff_matrix = np.array(self.bezier_coeff(t))
+        control_points_matrix = np.array(self.control_points)
+        sample_points = coeff_matrix.dot(control_points_matrix)
+        if image_size is not None:
+            sample_points[:, 0] = sample_points[:, 0] * image_size[-1]
+            sample_points[:, -1] = sample_points[:, -1] * image_size[0]
+
+        return sample_points
+
+    def get_middle_control_points(self, x, y):
+        dy = y[1:] - y[:-1]
+        dx = x[1:] - x[:-1]
+        dt = (dx ** 2 + dy ** 2) ** 0.5
+        t = dt / dt.sum()
+        t = np.hstack(([0], t))
+        t = t.cumsum()
+        data = np.column_stack((x, y))
+        Pseudoinverse = np.linalg.pinv(self.bezier_coeff(t))  # (9,4) -> (4,9)
+        control_points = Pseudoinverse.dot(data)  # (4,9)*(9,2) -> (4,2)
+        medi_ctp = control_points[:, :].flatten().tolist()
+
+        return medi_ctp
+
+
+class BezierSampler(torch.nn.Module):
+    # Fast Batch Bezier sampler
+    def __init__(self, num_sample_points):
+        super().__init__()
+        self.num_control_points = 4
+        self.num_sample_points = num_sample_points
+        self.control_points = []
+        self.bezier_coeff = self.get_bezier_coefficient()
+        self.bernstein_matrix = self.get_bernstein_matrix()
+
+    def get_bezier_coefficient(self):
+        Mtk = lambda n, t, k: t ** k * (1 - t) ** (n - k) * n_over_k(n, k)
+        BezierCoeff = lambda ts: [[Mtk(3, t, k) for k in range(4)] for t in ts]
+        return BezierCoeff
+
+    def get_bernstein_matrix(self):
+        t = torch.linspace(0, 1, self.num_sample_points)
+        c_matrix = torch.tensor(self.bezier_coeff(t))
+        return c_matrix  # (num_sample_points, 4)
+
+    def get_sample_points(self, control_points_matrix):
+        if control_points_matrix.numel() == 0:
+            return control_points_matrix  # Looks better than a torch.Tensor
+        if self.bernstein_matrix.device != control_points_matrix.device:
+            self.bernstein_matrix = self.bernstein_matrix.to(control_points_matrix.device)
+
+        return upcast(self.bernstein_matrix).matmul(upcast(control_points_matrix))
+
+
+@torch.no_grad()
+def get_valid_points(points):
+    # ... x 2
+    if points.numel() == 0:
+        return torch.tensor([1], dtype=torch.bool, device=points.device)
+    return (points[..., 0] > 0) * (points[..., 0] < 1) * (points[..., 1] > 0) * (points[..., 1] < 1)
+
+
+@torch.no_grad()
+def cubic_bezier_curve_segment(control_points, sample_points):
+    # Cut a batch of cubic bezier curves to its in-image segments (assume at least 2 valid sample points per curve).
+    # Based on De Casteljau's algorithm, formula for cubic bezier curve is derived by:
+    # https://stackoverflow.com/a/11704152/15449902
+    # control_points: B x 4 x 2
+    # sample_points: B x N x 2
+    if control_points.numel() == 0 or sample_points.numel() == 0:
+        return control_points
+    B, N = sample_points.shape[:-1]
+    valid_points = get_valid_points(sample_points)  # B x N, bool
+    t = torch.linspace(0.0, 1.0, steps=N, dtype=sample_points.dtype, device=sample_points.device)
+
+    # First & Last valid index (B)
+    # Get unique values for deterministic behaviour on cuda:
+    # https://pytorch.org/docs/1.6.0/generated/torch.max.html?highlight=max#torch.max
+    t0 = t[(valid_points + torch.arange(N, device=valid_points.device).flip([0]) * valid_points).max(dim=-1).indices]
+    t1 = t[(valid_points + torch.arange(N, device=valid_points.device) * valid_points).max(dim=-1).indices]
+
+    # Generate transform matrix (old control points -> new control points = linear transform)
+    u0 = 1 - t0  # B
+    u1 = 1 - t1  # B
+    transform_matrix_c = [torch.stack([u0 ** (3 - i) * u1 ** i for i in range(4)], dim=-1),
+                          torch.stack([3 * t0 * u0 ** 2,
+                                       2 * t0 * u0 * u1 + u0 ** 2 * t1,
+                                       t0 * u1 ** 2 + 2 * u0 * u1 * t1,
+                                       3 * t1 * u1 ** 2], dim=-1),
+                          torch.stack([3 * t0 ** 2 * u0,
+                                       t0 ** 2 * u1 + 2 * t0 * t1 * u0,
+                                       2 * t0 * t1 * u1 + t1 ** 2 * u0,
+                                       3 * t1 ** 2 * u1], dim=-1),
+                          torch.stack([t0 ** (3 - i) * t1 ** i for i in range(4)], dim=-1)]
+    transform_matrix = torch.stack(transform_matrix_c, dim=-2).transpose(-2, -1)  # B x 4 x 4, f**k this!
+    transform_matrix = transform_matrix.unsqueeze(1).expand(B, 2, 4, 4)
+
+    # Matrix multiplication
+    res = transform_matrix.matmul(control_points.permute(0, 2, 1).unsqueeze(-1))  # B x 2 x 4 x 1
+
+    return res.squeeze(-1).permute(0, 2, 1)

adet/utils/misc.py

+from typing import List, Optional
+import torch
+from torch.functional import Tensor
+from torchvision.ops.boxes import box_area
+import torch.distributed as dist
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+@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)]
+    if target.ndim == 2:
+        assert output.ndim == 3
+        output = output.mean(1)
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, -1)
+    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 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_xyxy_to_cxcywh(x):
+    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
+def box_iou(boxes1, boxes2):
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
+    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+
+    union = area1[:, None] + area2 - inter
+
+    iou = inter / union
+    return iou, union
+
+
+def generalized_box_iou(boxes1, boxes2):
+    """
+    Generalized IoU from https://giou.stanford.edu/
+    The boxes should be in [x0, y0, x1, y1] format
+    Returns a [N, M] pairwise matrix, where N = len(boxes1)
+    and M = len(boxes2)
+    """
+    # degenerate boxes gives inf / nan results
+    # so do an early check
+    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
+    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
+    iou, union = box_iou(boxes1, boxes2)
+
+    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
+    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,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 inverse_sigmoid(x, eps=1e-5):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1/x2)
+
+def sigmoid_offset(x, offset=True):
+    # modified sigmoid for range [-0.5, 1.5]
+    if offset:
+        return x.sigmoid() * 2 - 0.5
+    else:
+        return x.sigmoid()
+
+def inverse_sigmoid_offset(x, eps=1e-5, offset=True):
+    if offset:
+        x = (x + 0.5) / 2.0
+    return inverse_sigmoid(x, eps)
+
+def _max_by_axis(the_list):
+    # type: (List[List[int]]) -> List[int]
+    maxes = the_list[0]
+    for sublist in the_list[1:]:
+        for index, item in enumerate(sublist):
+            maxes[index] = max(maxes[index], item)
+    return maxes
+
+
+def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
+    # make this more general
+    if tensor_list[0].ndim == 3:
+        # make it support different-sized images
+        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
+        # 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)
+            m[: img.shape[1], :img.shape[2]] = False
+    else:
+        raise ValueError('not supported')
+    return NestedTensor(tensor, mask)
+
+
+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)

adet/utils/visualizer.py

+import numpy as np
+import pickle
+from detectron2.utils.visualizer import Visualizer
+import matplotlib.colors as mplc
+import matplotlib.font_manager as mfm
+import matplotlib as mpl
+import matplotlib.figure as mplfigure
+import random
+from shapely.geometry import LineString
+import math
+import operator
+from functools import reduce
+
+class TextVisualizer(Visualizer):
+    def __init__(self, image, metadata, instance_mode, cfg):
+        Visualizer.__init__(self, image, metadata, instance_mode=instance_mode)
+        self.voc_size = cfg.MODEL.TRANSFORMER.VOC_SIZE
+        self.use_customer_dictionary = cfg.MODEL.TRANSFORMER.CUSTOM_DICT
+        if self.voc_size == 96:
+            self.CTLABELS = [' ','!','"','#','$','%','&','\'','(',')','*','+',',','-','.','/','0','1','2','3','4','5','6','7','8','9',':',';','<','=','>','?','@','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z','[','\\',']','^','_','`','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z','{','|','}','~']
+        elif self.voc_size == 37:
+            self.CTLABELS = ['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z','0','1','2','3','4','5','6','7','8','9']
+        else:
+            with open(self.use_customer_dictionary, 'rb') as fp:
+                self.CTLABELS = pickle.load(fp)
+        # voc_size includes the unknown class, which is not in self.CTABLES
+        assert(int(self.voc_size - 1) == len(self.CTLABELS)), "voc_size is not matched dictionary size, got {} and {}.".format(int(self.voc_size - 1), len(self.CTLABELS))
+
+    def draw_instance_predictions(self, predictions):
+        ctrl_pnts = predictions.ctrl_points.numpy()
+        scores = predictions.scores.tolist()
+        recs = predictions.recs
+        bd_pts = np.asarray(predictions.bd)
+
+        self.overlay_instances(ctrl_pnts, scores, recs, bd_pts)
+
+        return self.output
+
+    def _process_ctrl_pnt(self, pnt):
+        points = pnt.reshape(-1, 2)
+        return points
+
+    def _ctc_decode_recognition(self, rec):
+        last_char = '###'
+        s = ''
+        for c in rec:
+            c = int(c)
+            if c < self.voc_size - 1:
+                if last_char != c:
+                    if self.voc_size == 37 or self.voc_size == 96:
+                        s += self.CTLABELS[c]
+                        last_char = c
+                    else:
+                        s += str(chr(self.CTLABELS[c]))
+                        last_char = c
+            else:
+                last_char = '###'
+        return s
+
+    def overlay_instances(self, ctrl_pnts, scores, recs, bd_pnts, alpha=0.4):
+        colors = [(0,0.5,0),(0,0.75,0),(1,0,1),(0.75,0,0.75),(0.5,0,0.5),(1,0,0),(0.75,0,0),(0.5,0,0),
+        (0,0,1),(0,0,0.75),(0.75,0.25,0.25),(0.75,0.5,0.5),(0,0.75,0.75),(0,0.5,0.5),(0,0.3,0.75)]
+
+        for ctrl_pnt, score, rec, bd in zip(ctrl_pnts, scores, recs, bd_pnts):
+            color = random.choice(colors)
+
+            # draw polygons
+            if bd is not None:
+                bd = np.hsplit(bd, 2)
+                bd = np.vstack([bd[0], bd[1][::-1]])
+                self.draw_polygon(bd, color, alpha=alpha)
+
+            # draw center lines
+            line = self._process_ctrl_pnt(ctrl_pnt)
+            line_ = LineString(line)
+            center_point = np.array(line_.interpolate(0.5, normalized=True).coords[0], dtype=np.int32)
+            # self.draw_line(
+            #     line[:, 0],
+            #     line[:, 1],
+            #     color=color,
+            #     linewidth=2
+            # )
+            # for pt in line:
+            #     self.draw_circle(pt, 'w', radius=4)
+            #     self.draw_circle(pt, 'r', radius=2)
+
+            # draw text
+            text = self._ctc_decode_recognition(rec)
+            if self.voc_size == 37:
+                text = text.upper()
+            # text = "{:.2f}: {}".format(score, text)
+            text = "{}".format(text)
+            lighter_color = self._change_color_brightness(color, brightness_factor=0)
+            if bd is not None:
+                text_pos = bd[0] - np.array([0,15])
+            else:
+                text_pos = center_point
+            horiz_align = "left"
+            font_size = self._default_font_size
+            self.draw_text(
+                        text,
+                        text_pos,
+                        color=lighter_color,
+                        horizontal_alignment=horiz_align,
+                        font_size=font_size,
+                        draw_chinese=False if self.voc_size == 37 or self.voc_size == 96 else True
+                    )
+
+    def draw_text(
+        self,
+        text,
+        position,
+        *,
+        font_size=None,
+        color="g",
+        horizontal_alignment="center",
+        rotation=0,
+        draw_chinese=False
+    ):
+        """
+        Args:
+            text (str): class label
+            position (tuple): a tuple of the x and y coordinates to place text on image.
+            font_size (int, optional): font of the text. If not provided, a font size
+                proportional to the image width is calculated and used.
+            color: color of the text. Refer to `matplotlib.colors` for full list
+                of formats that are accepted.
+            horizontal_alignment (str): see `matplotlib.text.Text`
+            rotation: rotation angle in degrees CCW
+        Returns:
+            output (VisImage): image object with text drawn.
+        """
+        if not font_size:
+            font_size = self._default_font_size
+
+        # since the text background is dark, we don't want the text to be dark
+        color = np.maximum(list(mplc.to_rgb(color)), 0.2)
+        color[np.argmax(color)] = max(0.8, np.max(color))
+        
+        x, y = position
+        if draw_chinese:
+            font_path = "./simsun.ttc"
+            prop = mfm.FontProperties(fname=font_path)
+            self.output.ax.text(
+                x,
+                y,
+                text,
+                size=font_size * self.output.scale,
+                family="sans-serif",
+                bbox={"facecolor": "white", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"},
+                verticalalignment="top",
+                horizontalalignment=horizontal_alignment,
+                color=color,
+                zorder=10,
+                rotation=rotation,
+                fontproperties=prop
+            )
+        else:
+            self.output.ax.text(
+                x,
+                y,
+                text,
+                size=font_size * self.output.scale,
+                family="sans-serif",
+                bbox={"facecolor": "white", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"},
+                verticalalignment="top",
+                horizontalalignment=horizontal_alignment,
+                color=color,
+                zorder=10,
+                rotation=rotation,
+            )
+        return self.output
\ No newline at end of file

configs/R_101/Base_det.yaml

+MODEL:
+  META_ARCHITECTURE: "TransformerPureDetector"
+  MASK_ON: False
+  PIXEL_MEAN: [123.675, 116.280, 103.530]
+  PIXEL_STD: [58.395, 57.120, 57.375]
+  BACKBONE:
+    NAME: "build_resnet_backbone"
+  RESNETS:
+    DEPTH: 101
+    STRIDE_IN_1X1: False
+    OUT_FEATURES: ["res3", "res4", "res5"]
+  TRANSFORMER:
+    ENABLED: True
+    NUM_FEATURE_LEVELS: 4
+    TEMPERATURE: 10000
+    ENC_LAYERS: 6
+    DEC_LAYERS: 6
+    DIM_FEEDFORWARD: 1024
+    HIDDEN_DIM: 256
+    DROPOUT: 0.0
+    NHEADS: 8
+    NUM_QUERIES: 100
+    ENC_N_POINTS: 4
+    DEC_N_POINTS: 4
+    NUM_POINTS: 25
+    INFERENCE_TH_TEST: 0.4
+    LOSS:
+      BEZIER_SAMPLE_POINTS: 25
+      BEZIER_CLASS_WEIGHT: 1.0
+      BEZIER_COORD_WEIGHT: 1.0
+      POINT_CLASS_WEIGHT: 1.0
+      POINT_COORD_WEIGHT: 1.0
+      POINT_TEXT_WEIGHT: 0.5
+      BOUNDARY_WEIGHT: 0.5
+
+SOLVER:
+  WEIGHT_DECAY: 1e-4
+  OPTIMIZER: "ADAMW"
+  LR_BACKBONE_NAMES: ['backbone.0']
+  LR_LINEAR_PROJ_NAMES: ['reference_points', 'sampling_offsets']
+  LR_LINEAR_PROJ_MULT: 1.
+  CLIP_GRADIENTS:
+    ENABLED: True
+    CLIP_TYPE: "full_model"
+    CLIP_VALUE: 0.1
+    NORM_TYPE: 2.0
+
+INPUT:
+  HFLIP_TRAIN: False
+  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800, 832, 864, 896)
+  MAX_SIZE_TRAIN: 1600
+  MIN_SIZE_TEST: 1000
+  MAX_SIZE_TEST: 1892
+  CROP:
+    ENABLED: True
+    CROP_INSTANCE: False
+    SIZE: [0.1, 0.1]
+  FORMAT: "RGB"
+
+DATALOADER:
+  NUM_WORKERS: 8
+
+VERSION: 2
+SEED: 42
\ No newline at end of file

configs/R_101/TotalText/finetune_150k_tt_mlt_13_15.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R101/150k_tt_mlt_13_15/pretrain/model_final.pth"
+
+DATASETS:
+  TRAIN: ("totaltext_train",)
+  TEST: ("totaltext_test",)
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (100000,)  # no step
+  MAX_ITER: 10000
+  CHECKPOINT_PERIOD: 2000
+
+TEST:
+  EVAL_PERIOD: 1000
+
+OUTPUT_DIR: "output/R101/150k_tt_mlt_13_15/finetune/totaltext"
\ No newline at end of file

configs/R_101/pretrain/150k_tt_mlt_13_15.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "pretrained_backbone/R-101.pkl"
+  
+DATASETS:
+  TRAIN: ("syntext1","syntext2","totaltext_train","mlt","ic13_train","ic15_train",)
+  TEST: ("totaltext_test",)
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-4
+  LR_BACKBONE: 1e-5
+  WARMUP_ITERS: 0
+  STEPS: (320000,)
+  MAX_ITER: 375000
+  CHECKPOINT_PERIOD: 100000
+
+TEST:
+  EVAL_PERIOD: 10000
+
+OUTPUT_DIR: "output/R101/150k_tt_mlt_13_15/pretrain"
\ No newline at end of file

configs/R_50/Base_det.yaml

+MODEL:
+  META_ARCHITECTURE: "TransformerPureDetector"
+  MASK_ON: False
+  PIXEL_MEAN: [123.675, 116.280, 103.530]
+  PIXEL_STD: [58.395, 57.120, 57.375]
+  BACKBONE:
+    NAME: "build_resnet_backbone"
+  RESNETS:
+    DEPTH: 50
+    STRIDE_IN_1X1: False
+    OUT_FEATURES: ["res3", "res4", "res5"]
+  TRANSFORMER:
+    ENABLED: True
+    NUM_FEATURE_LEVELS: 4
+    TEMPERATURE: 10000
+    ENC_LAYERS: 6
+    DEC_LAYERS: 6
+    DIM_FEEDFORWARD: 1024
+    HIDDEN_DIM: 256
+    DROPOUT: 0.0
+    NHEADS: 8
+    NUM_QUERIES: 100
+    ENC_N_POINTS: 4
+    DEC_N_POINTS: 4
+    NUM_POINTS: 25
+    INFERENCE_TH_TEST: 0.4
+    LOSS:
+      BEZIER_SAMPLE_POINTS: 25
+      BEZIER_CLASS_WEIGHT: 1.0
+      BEZIER_COORD_WEIGHT: 1.0
+      POINT_CLASS_WEIGHT: 1.0
+      POINT_COORD_WEIGHT: 1.0
+      POINT_TEXT_WEIGHT: 0.5
+      BOUNDARY_WEIGHT: 0.5
+
+SOLVER:
+  WEIGHT_DECAY: 1e-4
+  OPTIMIZER: "ADAMW"
+  LR_BACKBONE_NAMES: ['backbone.0']
+  LR_LINEAR_PROJ_NAMES: ['reference_points', 'sampling_offsets']
+  LR_LINEAR_PROJ_MULT: 1.
+  CLIP_GRADIENTS:
+    ENABLED: True
+    CLIP_TYPE: "full_model"
+    CLIP_VALUE: 0.1
+    NORM_TYPE: 2.0
+
+INPUT:
+  HFLIP_TRAIN: False
+  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800, 832, 864, 896)
+  MAX_SIZE_TRAIN: 1600
+  MIN_SIZE_TEST: 1000
+  MAX_SIZE_TEST: 1892
+  CROP:
+    ENABLED: True
+    CROP_INSTANCE: False
+    SIZE: [0.1, 0.1]
+  FORMAT: "RGB"
+
+DATALOADER:
+  NUM_WORKERS: 8
+
+VERSION: 2
+SEED: 42
\ No newline at end of file

configs/R_50/CTW1500/finetune_96voc_50maxlen.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R50/ctw1500/pretrain_150k-tt-mlt-ic13-15_maxlen50_96voc/model_final.pth"
+  TRANSFORMER:
+    VOC_SIZE: 96
+    NUM_POINTS: 50
+    LOSS:
+        BEZIER_SAMPLE_POINTS: 50
+        BEZIER_CLASS_WEIGHT: 1.0
+        BEZIER_COORD_WEIGHT: 0.5
+        POINT_CLASS_WEIGHT: 1.0
+        POINT_COORD_WEIGHT: 0.5
+        POINT_TEXT_WEIGHT: 1.0 #0.5
+        BOUNDARY_WEIGHT: 0.25
+  
+DATASETS:
+  TRAIN: ("ctw1500_train_96voc",)
+  TEST: ("ctw1500_test",)
+
+INPUT:
+  ROTATE: False
+  MIN_SIZE_TEST: 1000
+  MAX_SIZE_TEST: 1200
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 5e-5
+  LR_BACKBONE: 5e-6
+  WARMUP_ITERS: 0
+  STEPS: (8000,)
+  MAX_ITER: 12000
+  CHECKPOINT_PERIOD: 4000
+
+TEST:
+  EVAL_PERIOD: 1000
+
+OUTPUT_DIR: "output/R50/ctw1500/finetune_maxlen50_96voc"

configs/R_50/CTW1500/pretrain_96voc_50maxlen.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "detectron2://ImageNetPretrained/torchvision/R-50.pkl"
+  TRANSFORMER:
+    VOC_SIZE: 96
+    NUM_POINTS: 50
+    LOSS:
+        BEZIER_SAMPLE_POINTS: 50  # the same as NUM_POINTS
+        BEZIER_CLASS_WEIGHT: 1.0
+        BEZIER_COORD_WEIGHT: 0.5
+        POINT_CLASS_WEIGHT: 1.0
+        POINT_COORD_WEIGHT: 0.5
+        POINT_TEXT_WEIGHT: 0.5
+        BOUNDARY_WEIGHT: 0.25
+
+DATASETS:
+  TRAIN: ("syntext1_96voc", "syntext2_96voc", "totaltext_train_96voc", "mlt_96voc", "ic15_train_96voc", "ic13_train_96voc",)
+  TEST: ("ctw1500_test",)
+
+INPUT:
+  MIN_SIZE_TEST: 1000
+  MAX_SIZE_TEST: 1200
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-4
+  LR_BACKBONE: 1e-5
+  WARMUP_ITERS: 0
+  STEPS: (320000,)
+  MAX_ITER: 375000
+  CHECKPOINT_PERIOD: 100000
+
+TEST:
+  EVAL_PERIOD: 10000
+
+OUTPUT_DIR: "output/R50/ctw1500/pretrain_150k-tt-mlt-ic13-15_maxlen50_96voc"

configs/R_50/IC15/finetune_150k_tt_mlt_13_15.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R50/150k_tt_mlt_13_15/pretrain/model_final.pth"
+  TRANSFORMER:
+    INFERENCE_TH_TEST: 0.3
+
+DATASETS:
+  TRAIN: ("ic15_train",)
+  TEST: ("ic15_test",)
+
+INPUT:
+  MIN_SIZE_TRAIN: (800,900,1000,1100,1200,1300,1400)
+  MAX_SIZE_TRAIN: 3000
+  MIN_SIZE_TEST: 1440
+  MAX_SIZE_TEST: 4000
+  CROP:
+    ENABLED: False
+  ROTATE: False
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (100000,)  # no step
+  MAX_ITER: 3000
+  CHECKPOINT_PERIOD: 1000
+
+TEST:
+  EVAL_PERIOD: 1000
+  # 1 - Generic, 2 - Weak, 3 - Strong (for icdar2015)
+  LEXICON_TYPE: 3
+
+OUTPUT_DIR: "output/R50/150k_tt_mlt_13_15/finetune/ic15"
\ No newline at end of file

configs/R_50/IC15/finetune_150k_tt_mlt_13_15_textocr.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R50/150k_tt_mlt_13_15_textocr/pretrain/model_final.pth"
+  TRANSFORMER:
+    INFERENCE_TH_TEST: 0.3
+
+DATASETS:
+  TRAIN: ("ic15_train",)
+  TEST: ("ic15_test",)
+
+INPUT:
+  MIN_SIZE_TRAIN: (800,900,1000,1100,1200,1300,1400)
+  MAX_SIZE_TRAIN: 3000
+  MIN_SIZE_TEST: 1440
+  MAX_SIZE_TEST: 4000
+  CROP:
+    ENABLED: False
+  ROTATE: False
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (100000,)  # no step
+  MAX_ITER: 1000
+  CHECKPOINT_PERIOD: 1000
+
+TEST:
+  EVAL_PERIOD: 1000
+  # 1 - Generic, 2 - Weak, 3 - Strong (for icdar2015)
+  LEXICON_TYPE: 3
+
+OUTPUT_DIR: "output/R50/150k_tt_mlt_13_15_textocr/finetune/ic15"
\ No newline at end of file

configs/R_50/ReCTS/finetune.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "./output/R50/rects/pretrain/model_final.pth"
+  TRANSFORMER:
+    VOC_SIZE: 5462
+    CUSTOM_DICT: "chn_cls_list"
+    INFERENCE_TH_TEST: 0.35
+    LOSS:
+      POINT_TEXT_WEIGHT: 1.0
+  
+DATASETS:
+  TRAIN: ("rects_train", "rects_val",)
+  TEST: ("rects_test",)
+INPUT:
+  ROTATE: False
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (20000,)
+  MAX_ITER: 30000
+  CHECKPOINT_PERIOD: 30000
+
+TEST:
+  EVAL_PERIOD: 100000000
+
+OUTPUT_DIR: "output/R50/rects/finetune"
\ No newline at end of file

configs/R_50/ReCTS/pretrain.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "detectron2://ImageNetPretrained/torchvision/R-50.pkl"
+  TRANSFORMER:
+    VOC_SIZE: 5462
+  
+DATASETS:
+  TRAIN: ("chnsyn_train", "rects_train", "rects_val", "lsvt_train", "art_train",)
+  TEST: ("totaltext_test",)
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-4
+  LR_BACKBONE: 1e-5
+  WARMUP_ITERS: 0
+  STEPS: (300000,)
+  MAX_ITER: 400000
+  CHECKPOINT_PERIOD: 200000
+
+TEST:
+  EVAL_PERIOD: 100000000
+
+OUTPUT_DIR: "output/R50/rects/pretrain"

configs/R_50/TotalText/finetune_150k_tt.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R50/150k_tt/pretrain/model_final.pth"
+
+DATASETS:
+  TRAIN: ("totaltext_train",)
+  TEST: ("totaltext_test",)
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (100000,)  # no step
+  MAX_ITER: 10000
+  CHECKPOINT_PERIOD: 2000
+
+TEST:
+  EVAL_PERIOD: 1000
+
+OUTPUT_DIR: "output/R50/150k_tt/finetune/totaltext"
\ No newline at end of file

configs/R_50/TotalText/finetune_150k_tt_mlt_13_15.yaml

+_BASE_: "../Base_det.yaml"
+
+MODEL:
+  WEIGHTS: "output/R50/150k_tt_mlt_13_15/pretrain/model_final.pth"
+
+DATASETS:
+  TRAIN: ("totaltext_train",)
+  TEST: ("totaltext_test",)
+
+SOLVER:
+  IMS_PER_BATCH: 8
+  BASE_LR: 1e-5
+  LR_BACKBONE: 1e-6
+  WARMUP_ITERS: 0
+  STEPS: (100000,)  # no step
+  MAX_ITER: 10000
+  CHECKPOINT_PERIOD: 2000
+
+TEST:
+  EVAL_PERIOD: 1000
+
+OUTPUT_DIR: "output/R50/150k_tt_mlt_13_15/finetune/totaltext"
\ No newline at end of file