added textbpn++ detection module

IndicPhotoOCR/detection/textbpn/cfglib/config.py

@@ -0,0 +1,90 @@
+from easydict import EasyDict
+import torch
+import os
+
+config = EasyDict()
+
+
+# Normalize image
+config.means = (0.485, 0.456, 0.406)
+config.stds = (0.229, 0.224, 0.225)
+
+config.gpu = "1"
+
+# Experiment name #
+config.exp_name = "Synthtext"
+
+# dataloader jobs number
+config.num_workers = 24
+
+# batch_size
+config.batch_size = 12
+
+# training epoch number
+config.max_epoch = 200
+
+config.start_epoch = 0
+
+# learning rate
+config.lr = 1e-4
+
+# using GPU
+config.cuda = False
+
+config.output_dir = 'output'
+
+config.input_size = 640
+
+# max polygon per image
+# synText, total-text:64; CTW1500: 64; icdar: 64;  MLT: 32; TD500: 64.
+config.max_annotation = 64
+
+# adj num for graph
+config.adj_num = 4
+
+# control points number
+config.num_points = 20
+
+# use hard examples (annotated as '#')
+config.use_hard = True
+
+# Load data into memory at one time
+config.load_memory = False
+
+# prediction on 1/scale feature map
+config.scale = 1
+
+# # clip gradient of loss
+config.grad_clip = 25
+
+# demo tcl threshold
+config.dis_threshold = 0.4
+
+config.cls_threshold = 0.8
+
+# Contour approximation factor
+config.approx_factor = 0.004
+
+
+def update_config(config, extra_config):
+    for k, v in vars(extra_config).items():
+        config[k] = v
+    # print(config.gpu)
+    # config.device = torch.device('cuda') if config.cuda else torch.device('cpu')
+    config.device = torch.device('cpu')
+
+
+def print_config(config):
+    print('==========Options============')
+    for k, v in config.items():
+        print('{}: {}'.format(k, v))
+    print('=============End=============')
+
+
+
+################### MY Settings ##################
+config.resume=True
+
+config.device="cpu"
+
+# config.test_size = [224, 224]

IndicPhotoOCR/detection/textbpn/cfglib/option.py

@@ -0,0 +1,123 @@
+import argparse
+import torch
+import os
+import torch.backends.cudnn as cudnn
+
+from datetime import datetime
+
+
+def str2bool(v):
+    return v.lower() in ("yes", "true", "t", "1")
+
+
+def arg2str(args):
+    args_dict = vars(args)
+    option_str = datetime.now().strftime('%b%d_%H-%M-%S') + '\n'
+
+    for k, v in sorted(args_dict.items()):
+        option_str += ('{}: {}\n'.format(str(k), str(v)))
+
+    return option_str
+
+
+class BaseOptions(object):
+
+    def __init__(self):
+
+        self.parser = argparse.ArgumentParser()
+
+        # basic opts
+        self.parser.add_argument('--exp_name', default="TD500", type=str,
+                                 choices=['Synthtext', 'Totaltext', 'Ctw1500','Icdar2015',
+                                          "MLT2017", 'TD500', "MLT2019", "ArT", "ALL"], help='Experiment name')
+        self.parser.add_argument("--gpu", default="1", help="set gpu id", type=str)
+        self.parser.add_argument('--resume', default=None, type=str, help='Path to target resume checkpoint')
+        self.parser.add_argument('--num_workers', default=24, type=int, help='Number of workers used in dataloading')
+        self.parser.add_argument('--cuda', default=True, type=str2bool, help='Use cuda to train model')
+        self.parser.add_argument('--mgpu', action='store_true', help='Use multi-gpu to train model')
+        self.parser.add_argument('--save_dir', default='./model/', help='Path to save checkpoint models')
+        self.parser.add_argument('--vis_dir', default='./vis/', help='Path to save visualization images')
+        self.parser.add_argument('--log_dir', default='./logs/', help='Path to tensorboard log')
+        self.parser.add_argument('--loss', default='CrossEntropyLoss', type=str, help='Training Loss')
+        # self.parser.add_argument('--input_channel', default=1, type=int, help='number of input channels' )
+        self.parser.add_argument('--pretrain', default=False, type=str2bool, help='Pretrained AutoEncoder model')
+        self.parser.add_argument('--verbose', '-v', default=True, type=str2bool, help='Whether to output debug info')
+        self.parser.add_argument('--viz', action='store_true', help='Whether to output debug info')
+        # self.parser.add_argument('--viz', default=True, type=str2bool, help='Whether to output debug info')
+
+        # train opts
+        self.parser.add_argument('--max_epoch', default=250, type=int, help='Max epochs')
+        self.parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float, help='initial learning rate')
+        self.parser.add_argument('--lr_adjust', default='fix',
+                                 choices=['fix', 'poly'], type=str, help='Learning Rate Adjust Strategy')
+        self.parser.add_argument('--stepvalues', default=[], nargs='+', type=int, help='# of iter to change lr')
+        self.parser.add_argument('--weight_decay', '--wd', default=0., type=float, help='Weight decay for SGD')
+        self.parser.add_argument('--gamma', default=0.1, type=float, help='Gamma update for SGD lr')
+        self.parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
+        self.parser.add_argument('--batch_size', default=6, type=int, help='Batch size for training')
+        self.parser.add_argument('--optim', default='Adam', type=str, choices=['SGD', 'Adam'], help='Optimizer')
+        self.parser.add_argument('--save_freq', default=5, type=int, help='save weights every # epoch')
+        self.parser.add_argument('--display_freq', default=10, type=int, help='display training metrics every # iter')
+        self.parser.add_argument('--viz_freq', default=50, type=int, help='visualize training process every # iter')
+        self.parser.add_argument('--log_freq', default=10000, type=int, help='log to tensorboard every # iterations')
+        self.parser.add_argument('--val_freq', default=1000, type=int, help='do validation every # iterations')
+
+        # backbone
+        self.parser.add_argument('--scale', default=1, type=int, help='prediction on 1/scale feature map')
+        self.parser.add_argument('--net', default='resnet50', type=str,
+                                 choices=['vgg', 'resnet50', 'resnet18',
+                                          "deformable_resnet18", "deformable_resnet50"],
+                                 help='Network architecture')
+        # data args
+        self.parser.add_argument('--load_memory', default=False, type=str2bool, help='Load data into memory')
+        self.parser.add_argument('--rescale', type=float, default=255.0, help='rescale factor')
+        self.parser.add_argument('--input_size', default=640, type=int, help='model input size')
+        self.parser.add_argument('--test_size', default=[640, 960], type=int, nargs='+', help='test size')
+
+        # eval args00
+        self.parser.add_argument('--checkepoch', default=1070, type=int, help='Load checkpoint number')
+        self.parser.add_argument('--start_epoch', default=0, type=int, help='start epoch number')
+        self.parser.add_argument('--cls_threshold', default=0.875, type=float, help='threshold of pse')
+        self.parser.add_argument('--dis_threshold', default=0.35, type=float, help='filter the socre < score_i')
+
+        # demo args
+        self.parser.add_argument('--img_root', default=None,   type=str, help='Path to deploy images')
+
+    def parse(self, fixed=None):
+
+        if fixed is not None:
+            args = self.parser.parse_args(fixed)
+        else:
+            args = self.parser.parse_args()
+
+        return args
+
+    def initialize(self, fixed=None):
+
+        # Parse options
+        self.args = self.parse(fixed)
+        os.environ['CUDA_VISIBLE_DEVICES'] = self.args.gpu
+
+        # Setting default torch Tensor type
+        if self.args.cuda and torch.cuda.is_available():
+            torch.set_default_tensor_type('torch.cuda.FloatTensor')
+            cudnn.benchmark = True
+        else:
+            torch.set_default_tensor_type('torch.FloatTensor')
+
+        # Create weights saving directory
+        if not os.path.exists(self.args.save_dir):
+            os.mkdir(self.args.save_dir)
+
+        # Create weights saving directory of target model
+        model_save_path = os.path.join(self.args.save_dir, self.args.exp_name)
+
+        if not os.path.exists(model_save_path):
+            os.mkdir(model_save_path)
+
+        return self.args
+
+    def update(self, args, extra_options):
+
+        for k, v in extra_options.items():
+            setattr(args, k, v)

IndicPhotoOCR/detection/textbpn/models/TextBPN_resnet50_300.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b735b9c93c8758972d3b8cfd3ef8e1c09afa8cd9106f4cb11406300b141b1d78
+size 145703602

IndicPhotoOCR/detection/textbpn/network/Reg_loss.py

@@ -0,0 +1,196 @@
+# -*- coding: utf-8 -*-
+# @Time    : 10/1/21
+# @Author  : GXYM
+import torch
+from torch import nn
+import numpy as np
+import torch.nn.functional as F
+
+
+class PolyMatchingLoss(nn.Module):
+    def __init__(self, pnum, device, loss_type="L1"):
+        super(PolyMatchingLoss, self).__init__()
+
+        self.pnum = pnum
+        self.device = device
+        self.loss_type = loss_type
+        self.smooth_L1 = F.smooth_l1_loss
+        self.L2_loss = torch.nn.MSELoss(reduce=False, size_average=False)
+
+        batch_size = 1
+        pidxall = np.zeros(shape=(batch_size, pnum, pnum), dtype=np.int32)
+        for b in range(batch_size):
+            for i in range(pnum):
+                pidx = (np.arange(pnum) + i) % pnum
+                pidxall[b, i] = pidx
+
+        pidxall = torch.from_numpy(np.reshape(pidxall, newshape=(batch_size, -1))).to(device)
+        self.feature_id = pidxall.unsqueeze_(2).long().expand(pidxall.size(0), pidxall.size(1), 2).detach()
+        print(self.feature_id.shape)
+
+    def match_loss(self, pred, gt):
+        batch_size = pred.shape[0]
+        feature_id = self.feature_id.expand(batch_size, self.feature_id.size(1), 2)
+
+        gt_expand = torch.gather(gt, 1, feature_id).view(batch_size, self.pnum, self.pnum, 2)
+        pred_expand = pred.unsqueeze(1)
+
+        if self.loss_type == "L2":
+            dis = self.L2_loss(pred_expand, gt_expand)
+            dis = dis.sum(3).sqrt().mean(2)
+        elif self.loss_type == "L1":
+            dis = self.smooth_L1(pred_expand, gt_expand, reduction='none')
+            dis = dis.sum(3).mean(2)
+
+        min_dis, min_id = torch.min(dis, dim=1, keepdim=True)
+
+        return min_dis
+
+    def forward(self, pred_list, gt):
+        loss = torch.tensor(0.)
+        for pred in pred_list:
+            loss += torch.mean(self.match_loss(pred, gt))
+
+        return loss / torch.tensor(len(pred_list))
+
+        # los = []
+        # for pred in pred_list:
+        #     los.append(self.match_loss(pred, gt))
+        #
+        # los_b = torch.tensor(0.)
+        # loss_c = torch.tensor(0.)
+        # for i, _ in enumerate(los):
+        #     los_b += torch.mean(los[i])
+        #     loss_c += (torch.mean(torch.clamp(los[i] - los[i - 1], min=0.0)) if i > 0 else torch.tensor(0.))
+        # loss = los_b / torch.tensor(len(los)) + 0.5*loss_c / torch.tensor(len(los)-1)
+        #
+        # return loss
+
+
+class AttentionLoss(nn.Module):
+    def __init__(self, beta=4, gamma=0.5):
+        super(AttentionLoss, self).__init__()
+
+        self.beta = beta
+        self.gamma = gamma
+
+    def forward(self, pred, gt):
+        num_pos = torch.sum(gt)
+        num_neg = torch.sum(1 - gt)
+        alpha = num_neg / (num_pos + num_neg)
+        edge_beta = torch.pow(self.beta, torch.pow(1 - pred, self.gamma))
+        bg_beta = torch.pow(self.beta, torch.pow(pred, self.gamma))
+
+        loss = 0
+        loss = loss - alpha * edge_beta * torch.log(pred) * gt
+        loss = loss - (1 - alpha) * bg_beta * torch.log(1 - pred) * (1 - gt)
+        return torch.mean(loss)
+
+
+class GeoCrossEntropyLoss(nn.Module):
+    def __init__(self):
+        super(GeoCrossEntropyLoss, self).__init__()
+
+    def forward(self, output, target, poly):
+        output = torch.nn.functional.softmax(output, dim=1)
+        output = torch.log(torch.clamp(output, min=1e-4))
+        poly = poly.view(poly.size(0), 4, poly.size(1) // 4, 2)
+        target = target[..., None, None].expand(poly.size(0), poly.size(1), 1, poly.size(3))
+        target_poly = torch.gather(poly, 2, target)
+        sigma = (poly[:, :, 0] - poly[:, :, 1]).pow(2).sum(2, keepdim=True)
+        kernel = torch.exp(-(poly - target_poly).pow(2).sum(3) / (sigma / 3))
+        loss = -(output * kernel.transpose(2, 1)).sum(1).mean()
+        return loss
+
+
+class AELoss(nn.Module):
+    def __init__(self):
+        super(AELoss, self).__init__()
+
+    def forward(self, ae, ind, ind_mask):
+        """
+        ae: [b, 1, h, w]
+        ind: [b, max_objs, max_parts]
+        ind_mask: [b, max_objs, max_parts]
+        obj_mask: [b, max_objs]
+        """
+        # first index
+        b, _, h, w = ae.shape
+        b, max_objs, max_parts = ind.shape
+        obj_mask = torch.sum(ind_mask, dim=2) != 0
+
+        ae = ae.view(b, h * w, 1)
+        seed_ind = ind.view(b, max_objs * max_parts, 1)
+        tag = ae.gather(1, seed_ind).view(b, max_objs, max_parts)
+
+        # compute the mean
+        tag_mean = tag * ind_mask
+        tag_mean = tag_mean.sum(2) / (ind_mask.sum(2) + 1e-4)
+
+        # pull ae of the same object to their mean
+        pull_dist = (tag - tag_mean.unsqueeze(2)).pow(2) * ind_mask
+        obj_num = obj_mask.sum(dim=1).float()
+        pull = (pull_dist.sum(dim=(1, 2)) / (obj_num + 1e-4)).sum()
+        pull /= b
+
+        # push away the mean of different objects
+        push_dist = torch.abs(tag_mean.unsqueeze(1) - tag_mean.unsqueeze(2))
+        push_dist = 1 - push_dist
+        push_dist = nn.functional.relu(push_dist, inplace=True)
+        obj_mask = (obj_mask.unsqueeze(1) + obj_mask.unsqueeze(2)) == 2
+        push_dist = push_dist * obj_mask.float()
+        push = ((push_dist.sum(dim=(1, 2)) - obj_num) / (obj_num * (obj_num - 1) + 1e-4)).sum()
+        push /= b
+        return pull, push
+
+
+def smooth_l1_loss(inputs, target, sigma=9.0):
+    try:
+        diff = torch.abs(inputs - target)
+        less_one = (diff < 1.0 / sigma).float()
+        loss = less_one * 0.5 * diff ** 2 * sigma \
+               + torch.abs(torch.tensor(1.0) - less_one) * (diff - 0.5 / sigma)
+        loss = torch.mean(loss) if loss.numel() > 0 else torch.tensor(0.0)
+    except Exception as e:
+        print('RPN_REGR_Loss Exception:', e)
+        loss = torch.tensor(0.0)
+
+    return loss
+
+
+def _neg_loss(pred, gt):
+    ''' Modified focal loss. Exactly the same as CornerNet.
+        Runs faster and costs a little bit more memory
+        Arguments:
+            pred (batch x c x h x w)
+            gt_regr (batch x c x h x w)
+    '''
+    pos_inds = gt.eq(1).float()
+    neg_inds = gt.lt(1).float()
+
+    neg_weights = torch.pow(1 - gt, 4)
+
+    loss = 0
+
+    pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
+    neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds
+
+    num_pos = pos_inds.float().sum()
+    pos_loss = pos_loss.sum()
+    neg_loss = neg_loss.sum()
+
+    if num_pos == 0:
+        loss = loss - neg_loss
+    else:
+        loss = loss - (pos_loss + neg_loss) / num_pos
+    return loss
+
+
+class FocalLoss(nn.Module):
+    '''nn.Module warpper for focal loss'''
+    def __init__(self):
+        super(FocalLoss, self).__init__()
+        self.neg_loss = _neg_loss
+
+    def forward(self, out, target):
+        return self.neg_loss(out, target)

IndicPhotoOCR/detection/textbpn/network/Seg_loss.py

@@ -0,0 +1,107 @@
+# -*- coding: utf-8 -*-
+# @Time    : 10/1/21
+# @Author  : GXYM
+import torch
+from torch import nn
+import numpy as np
+
+
+class SegmentLoss(nn.Module):
+    def __init__(self, Lambda, ratio=3, reduction='mean'):
+        """Implement PSE Loss.
+        """
+        super(SegmentLoss, self).__init__()
+        assert reduction in ['mean', 'sum'], " reduction must in ['mean','sum']"
+        self.Lambda = Lambda
+        self.ratio = ratio
+        self.reduction = reduction
+
+    def forward(self, outputs, labels, training_masks, th=0.5):
+        texts = outputs[:, -1, :, :]
+        kernels = outputs[:, :-1, :, :]
+        gt_texts = labels[:, -1, :, :]
+        gt_kernels = labels[:, :-1, :, :]
+
+        selected_masks = self.ohem_batch(texts, gt_texts, training_masks)
+        selected_masks = selected_masks.to(outputs.device)
+
+        loss_text = self.dice_loss(texts, gt_texts, selected_masks)
+
+        loss_kernels = []
+        # mask0 = torch.sigmoid(texts).data.cpu().numpy()
+        mask0 = texts.data.cpu().numpy()
+        mask1 = training_masks.data.cpu().numpy()
+        selected_masks = ((mask0 > th) & (mask1 > th)).astype('float32')
+        selected_masks = torch.from_numpy(selected_masks).float()
+        selected_masks = selected_masks.to(outputs.device)
+        kernels_num = gt_kernels.size()[1]
+        for i in range(kernels_num):
+            kernel_i = kernels[:, i, :, :]
+            gt_kernel_i = gt_kernels[:, i, :, :]
+            loss_kernel_i = self.dice_loss(kernel_i, gt_kernel_i, selected_masks)
+            loss_kernels.append(loss_kernel_i)
+        loss_kernels = torch.stack(loss_kernels).mean(0)
+        if self.reduction == 'mean':
+            loss_text = loss_text.mean()
+            loss_kernels = loss_kernels.mean()
+        elif self.reduction == 'sum':
+            loss_text = loss_text.sum()
+            loss_kernels = loss_kernels.sum()
+
+        loss = self.Lambda *loss_text + (1-self.Lambda)*loss_kernels
+        return loss_text, loss_kernels, loss
+
+    def dice_loss(self, input, target, mask):
+        # input = torch.sigmoid(input)
+
+        input = input.contiguous().view(input.size()[0], -1)
+        target = target.contiguous().view(target.size()[0], -1)
+        mask = mask.contiguous().view(mask.size()[0], -1)
+
+        input = input * mask
+        target = (target.float()) * mask
+
+        a = torch.sum(input * target, 1)
+        b = torch.sum(input * input, 1) + 0.001
+        c = torch.sum(target * target, 1) + 0.001
+        d = (2 * a) / (b + c)
+        return 1 - d
+
+    def ohem_single(self, score, gt_text, training_mask, th=0.5):
+        pos_num = (int)(np.sum(gt_text > th)) - (int)(np.sum((gt_text > th) & (training_mask <= th)))
+
+        if pos_num == 0:
+            # selected_mask = gt_text.copy() * 0 # may be not good
+            selected_mask = training_mask
+            selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
+            return selected_mask
+
+        neg_num = (int)(np.sum(gt_text <= th))
+        neg_num = (int)(min(pos_num * 3, neg_num))
+
+        if neg_num == 0:
+            selected_mask = training_mask
+            selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
+            return selected_mask
+
+        neg_score = score[gt_text <= th]
+        # 将负样本得分从高到低排序
+        neg_score_sorted = np.sort(-neg_score)
+        threshold = -neg_score_sorted[neg_num - 1]
+        # 选出 得分高的 负样本 和正样本 的 mask
+        selected_mask = ((score >= threshold) | (gt_text > th)) & (training_mask > th)
+        selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
+        return selected_mask
+
+    def ohem_batch(self, scores, gt_texts, training_masks):
+        scores = scores.data.cpu().numpy()
+        gt_texts = gt_texts.data.cpu().numpy()
+        training_masks = training_masks.data.cpu().numpy()
+        selected_masks = []
+        for i in range(scores.shape[0]):
+            selected_masks.append(self.ohem_single(scores[i, :, :], gt_texts[i, :, :], training_masks[i, :, :]))
+
+        selected_masks = np.concatenate(selected_masks, 0)
+        selected_masks = torch.from_numpy(selected_masks).float()
+
+        return selected_masks

IndicPhotoOCR/detection/textbpn/network/__init__.py

@@ -0,0 +1 @@
+from . import *

IndicPhotoOCR/detection/textbpn/network/backbone/__init__.py

@@ -0,0 +1 @@
+from .resnet import resnet18, resnet34, resnet50, resnet101, deformable_resnet50, deformable_resnet18

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/Makefile

@@ -0,0 +1,6 @@
+#!/bin/bash
+rm  *.so 
+python setup.py build_ext --inplace
+rm -rf ./build
+
+

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/Makefile.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+rm  *.so 
+python setup.py build_ext --inplace
+rm -rf ./build
+
+

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/__init__.py

@@ -0,0 +1,13 @@
+from .functions.deform_conv import deform_conv, modulated_deform_conv
+from .functions.deform_pool import deform_roi_pooling
+from .modules.deform_conv import (DeformConv, ModulatedDeformConv,
+                                  DeformConvPack, ModulatedDeformConvPack)
+from .modules.deform_pool import (DeformRoIPooling, DeformRoIPoolingPack,
+                                  ModulatedDeformRoIPoolingPack)
+
+__all__ = [
+    'DeformConv', 'DeformConvPack', 'ModulatedDeformConv',
+    'ModulatedDeformConvPack', 'DeformRoIPooling', 'DeformRoIPoolingPack',
+    'ModulatedDeformRoIPoolingPack', 'deform_conv', 'modulated_deform_conv',
+    'deform_roi_pooling'
+]

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/functions/deform_conv.py

@@ -0,0 +1,181 @@
+import torch
+from torch.autograd import Function
+from torch.nn.modules.utils import _pair
+
+from .. import deform_conv_cuda
+
+
+class DeformConvFunction(Function):
+
+    @staticmethod
+    def forward(ctx,
+                input,
+                offset,
+                weight,
+                stride=1,
+                padding=0,
+                dilation=1,
+                groups=1,
+                deformable_groups=1,
+                im2col_step=64):
+        if input is not None and input.dim() != 4:
+            raise ValueError(
+                "Expected 4D tensor as input, got {}D tensor instead.".format(
+                    input.dim()))
+        ctx.stride = _pair(stride)
+        ctx.padding = _pair(padding)
+        ctx.dilation = _pair(dilation)
+        ctx.groups = groups
+        ctx.deformable_groups = deformable_groups
+        ctx.im2col_step = im2col_step
+
+        ctx.save_for_backward(input, offset, weight)
+
+        output = input.new_empty(
+            DeformConvFunction._output_size(input, weight, ctx.padding,
+                                            ctx.dilation, ctx.stride))
+
+        ctx.bufs_ = [input.new_empty(0), input.new_empty(0)]  # columns, ones
+
+        if not input.is_cuda:
+            raise NotImplementedError
+        else:
+            cur_im2col_step = min(ctx.im2col_step, input.shape[0])
+            assert (input.shape[0] %
+                    cur_im2col_step) == 0, 'im2col step must divide batchsize'
+            deform_conv_cuda.deform_conv_forward_cuda(
+                input, weight, offset, output, ctx.bufs_[0], ctx.bufs_[1],
+                weight.size(3), weight.size(2), ctx.stride[1], ctx.stride[0],
+                ctx.padding[1], ctx.padding[0], ctx.dilation[1],
+                ctx.dilation[0], ctx.groups, ctx.deformable_groups,
+                cur_im2col_step)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, offset, weight = ctx.saved_tensors
+
+        grad_input = grad_offset = grad_weight = None
+
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+        else:
+            cur_im2col_step = min(ctx.im2col_step, input.shape[0])
+            assert (input.shape[0] %
+                    cur_im2col_step) == 0, 'im2col step must divide batchsize'
+
+            if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
+                grad_input = torch.zeros_like(input)
+                grad_offset = torch.zeros_like(offset)
+                deform_conv_cuda.deform_conv_backward_input_cuda(
+                    input, offset, grad_output, grad_input,
+                    grad_offset, weight, ctx.bufs_[0], weight.size(3),
+                    weight.size(2), ctx.stride[1], ctx.stride[0],
+                    ctx.padding[1], ctx.padding[0], ctx.dilation[1],
+                    ctx.dilation[0], ctx.groups, ctx.deformable_groups,
+                    cur_im2col_step)
+
+            if ctx.needs_input_grad[2]:
+                grad_weight = torch.zeros_like(weight)
+                deform_conv_cuda.deform_conv_backward_parameters_cuda(
+                    input, offset, grad_output,
+                    grad_weight, ctx.bufs_[0], ctx.bufs_[1], weight.size(3),
+                    weight.size(2), ctx.stride[1], ctx.stride[0],
+                    ctx.padding[1], ctx.padding[0], ctx.dilation[1],
+                    ctx.dilation[0], ctx.groups, ctx.deformable_groups, 1,
+                    cur_im2col_step)
+
+        return (grad_input, grad_offset, grad_weight, None, None, None, None,
+                None)
+
+    @staticmethod
+    def _output_size(input, weight, padding, dilation, stride):
+        channels = weight.size(0)
+        output_size = (input.size(0), channels)
+        for d in range(input.dim() - 2):
+            in_size = input.size(d + 2)
+            pad = padding[d]
+            kernel = dilation[d] * (weight.size(d + 2) - 1) + 1
+            stride_ = stride[d]
+            output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1, )
+        if not all(map(lambda s: s > 0, output_size)):
+            raise ValueError(
+                "convolution input is too small (output would be {})".format(
+                    'x'.join(map(str, output_size))))
+        return output_size
+
+
+class ModulatedDeformConvFunction(Function):
+
+    @staticmethod
+    def forward(ctx,
+                input,
+                offset,
+                mask,
+                weight,
+                bias=None,
+                stride=1,
+                padding=0,
+                dilation=1,
+                groups=1,
+                deformable_groups=1):
+        ctx.stride = stride
+        ctx.padding = padding
+        ctx.dilation = dilation
+        ctx.groups = groups
+        ctx.deformable_groups = deformable_groups
+        ctx.with_bias = bias is not None
+        if not ctx.with_bias:
+            bias = input.new_empty(1)  # fake tensor
+        if not input.is_cuda:
+            raise NotImplementedError
+        if weight.requires_grad or mask.requires_grad or offset.requires_grad \
+                or input.requires_grad:
+            ctx.save_for_backward(input, offset, mask, weight, bias)
+        output = input.new_empty(
+            ModulatedDeformConvFunction._infer_shape(ctx, input, weight))
+        ctx._bufs = [input.new_empty(0), input.new_empty(0)]
+        deform_conv_cuda.modulated_deform_conv_cuda_forward(
+            input, weight, bias, ctx._bufs[0], offset, mask, output,
+            ctx._bufs[1], weight.shape[2], weight.shape[3], ctx.stride,
+            ctx.stride, ctx.padding, ctx.padding, ctx.dilation, ctx.dilation,
+            ctx.groups, ctx.deformable_groups, ctx.with_bias)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+        input, offset, mask, weight, bias = ctx.saved_tensors
+        grad_input = torch.zeros_like(input)
+        grad_offset = torch.zeros_like(offset)
+        grad_mask = torch.zeros_like(mask)
+        grad_weight = torch.zeros_like(weight)
+        grad_bias = torch.zeros_like(bias)
+        deform_conv_cuda.modulated_deform_conv_cuda_backward(
+            input, weight, bias, ctx._bufs[0], offset, mask, ctx._bufs[1],
+            grad_input, grad_weight, grad_bias, grad_offset, grad_mask,
+            grad_output, weight.shape[2], weight.shape[3], ctx.stride,
+            ctx.stride, ctx.padding, ctx.padding, ctx.dilation, ctx.dilation,
+            ctx.groups, ctx.deformable_groups, ctx.with_bias)
+        if not ctx.with_bias:
+            grad_bias = None
+
+        return (grad_input, grad_offset, grad_mask, grad_weight, grad_bias,
+                None, None, None, None, None)
+
+    @staticmethod
+    def _infer_shape(ctx, input, weight):
+        n = input.size(0)
+        channels_out = weight.size(0)
+        height, width = input.shape[2:4]
+        kernel_h, kernel_w = weight.shape[2:4]
+        height_out = (height + 2 * ctx.padding -
+                      (ctx.dilation * (kernel_h - 1) + 1)) // ctx.stride + 1
+        width_out = (width + 2 * ctx.padding -
+                     (ctx.dilation * (kernel_w - 1) + 1)) // ctx.stride + 1
+        return n, channels_out, height_out, width_out
+
+
+deform_conv = DeformConvFunction.apply
+modulated_deform_conv = ModulatedDeformConvFunction.apply

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/functions/deform_pool.py

@@ -0,0 +1,69 @@
+import torch
+from torch.autograd import Function
+
+from .. import deform_pool_cuda
+
+
+class DeformRoIPoolingFunction(Function):
+
+    @staticmethod
+    def forward(ctx,
+                data,
+                rois,
+                offset,
+                spatial_scale,
+                out_size,
+                out_channels,
+                no_trans,
+                group_size=1,
+                part_size=None,
+                sample_per_part=4,
+                trans_std=.0):
+        ctx.spatial_scale = spatial_scale
+        ctx.out_size = out_size
+        ctx.out_channels = out_channels
+        ctx.no_trans = no_trans
+        ctx.group_size = group_size
+        ctx.part_size = out_size if part_size is None else part_size
+        ctx.sample_per_part = sample_per_part
+        ctx.trans_std = trans_std
+
+        assert 0.0 <= ctx.trans_std <= 1.0
+        if not data.is_cuda:
+            raise NotImplementedError
+
+        n = rois.shape[0]
+        output = data.new_empty(n, out_channels, out_size, out_size)
+        output_count = data.new_empty(n, out_channels, out_size, out_size)
+        deform_pool_cuda.deform_psroi_pooling_cuda_forward(
+            data, rois, offset, output, output_count, ctx.no_trans,
+            ctx.spatial_scale, ctx.out_channels, ctx.group_size, ctx.out_size,
+            ctx.part_size, ctx.sample_per_part, ctx.trans_std)
+
+        if data.requires_grad or rois.requires_grad or offset.requires_grad:
+            ctx.save_for_backward(data, rois, offset)
+        ctx.output_count = output_count
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+
+        data, rois, offset = ctx.saved_tensors
+        output_count = ctx.output_count
+        grad_input = torch.zeros_like(data)
+        grad_rois = None
+        grad_offset = torch.zeros_like(offset)
+
+        deform_pool_cuda.deform_psroi_pooling_cuda_backward(
+            grad_output, data, rois, offset, output_count, grad_input,
+            grad_offset, ctx.no_trans, ctx.spatial_scale, ctx.out_channels,
+            ctx.group_size, ctx.out_size, ctx.part_size, ctx.sample_per_part,
+            ctx.trans_std)
+        return (grad_input, grad_rois, grad_offset, None, None, None, None,
+                None, None, None, None)
+
+
+deform_roi_pooling = DeformRoIPoolingFunction.apply

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/modules/deform_conv.py

@@ -0,0 +1,157 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.utils import _pair
+
+from ..functions.deform_conv import deform_conv, modulated_deform_conv
+
+
+class DeformConv(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 deformable_groups=1,
+                 bias=False):
+        super(DeformConv, self).__init__()
+
+        assert not bias
+        assert in_channels % groups == 0, \
+            'in_channels {} cannot be divisible by groups {}'.format(
+                in_channels, groups)
+        assert out_channels % groups == 0, \
+            'out_channels {} cannot be divisible by groups {}'.format(
+                out_channels, groups)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.padding = _pair(padding)
+        self.dilation = _pair(dilation)
+        self.groups = groups
+        self.deformable_groups = deformable_groups
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels // self.groups,
+                         *self.kernel_size))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, x, offset):
+        return deform_conv(x, offset, self.weight, self.stride, self.padding,
+                           self.dilation, self.groups, self.deformable_groups)
+
+
+class DeformConvPack(DeformConv):
+
+    def __init__(self, *args, **kwargs):
+        super(DeformConvPack, self).__init__(*args, **kwargs)
+
+        self.conv_offset = nn.Conv2d(
+            self.in_channels,
+            self.deformable_groups * 2 * self.kernel_size[0] *
+            self.kernel_size[1],
+            kernel_size=self.kernel_size,
+            stride=_pair(self.stride),
+            padding=_pair(self.padding),
+            bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset.weight.data.zero_()
+        self.conv_offset.bias.data.zero_()
+
+    def forward(self, x):
+        offset = self.conv_offset(x)
+        return deform_conv(x, offset, self.weight, self.stride, self.padding,
+                           self.dilation, self.groups, self.deformable_groups)
+
+
+class ModulatedDeformConv(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 deformable_groups=1,
+                 bias=True):
+        super(ModulatedDeformConv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.deformable_groups = deformable_groups
+        self.with_bias = bias
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels // groups,
+                         *self.kernel_size))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.zero_()
+
+    def forward(self, x, offset, mask):
+        return modulated_deform_conv(x, offset, mask, self.weight, self.bias,
+                                     self.stride, self.padding, self.dilation,
+                                     self.groups, self.deformable_groups)
+
+
+class ModulatedDeformConvPack(ModulatedDeformConv):
+
+    def __init__(self, *args, **kwargs):
+        super(ModulatedDeformConvPack, self).__init__(*args, **kwargs)
+
+        self.conv_offset_mask = nn.Conv2d(
+            self.in_channels,
+            self.deformable_groups * 3 * self.kernel_size[0] *
+            self.kernel_size[1],
+            kernel_size=self.kernel_size,
+            stride=_pair(self.stride),
+            padding=_pair(self.padding),
+            bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, x):
+        out = self.conv_offset_mask(x)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset = torch.cat((o1, o2), dim=1)
+        mask = torch.sigmoid(mask)
+        return modulated_deform_conv(x, offset, mask, self.weight, self.bias,
+                                     self.stride, self.padding, self.dilation,
+                                     self.groups, self.deformable_groups)

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/modules/deform_pool.py

@@ -0,0 +1,172 @@
+from torch import nn
+
+from ..functions.deform_pool import deform_roi_pooling
+
+
+class DeformRoIPooling(nn.Module):
+
+    def __init__(self,
+                 spatial_scale,
+                 out_size,
+                 out_channels,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0):
+        super(DeformRoIPooling, self).__init__()
+        self.spatial_scale = spatial_scale
+        self.out_size = out_size
+        self.out_channels = out_channels
+        self.no_trans = no_trans
+        self.group_size = group_size
+        self.part_size = out_size if part_size is None else part_size
+        self.sample_per_part = sample_per_part
+        self.trans_std = trans_std
+
+    def forward(self, data, rois, offset):
+        if self.no_trans:
+            offset = data.new_empty(0)
+        return deform_roi_pooling(
+            data, rois, offset, self.spatial_scale, self.out_size,
+            self.out_channels, self.no_trans, self.group_size, self.part_size,
+            self.sample_per_part, self.trans_std)
+
+
+class DeformRoIPoolingPack(DeformRoIPooling):
+
+    def __init__(self,
+                 spatial_scale,
+                 out_size,
+                 out_channels,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0,
+                 num_offset_fcs=3,
+                 deform_fc_channels=1024):
+        super(DeformRoIPoolingPack,
+              self).__init__(spatial_scale, out_size, out_channels, no_trans,
+                             group_size, part_size, sample_per_part, trans_std)
+
+        self.num_offset_fcs = num_offset_fcs
+        self.deform_fc_channels = deform_fc_channels
+
+        if not no_trans:
+            seq = []
+            ic = self.out_size * self.out_size * self.out_channels
+            for i in range(self.num_offset_fcs):
+                if i < self.num_offset_fcs - 1:
+                    oc = self.deform_fc_channels
+                else:
+                    oc = self.out_size * self.out_size * 2
+                seq.append(nn.Linear(ic, oc))
+                ic = oc
+                if i < self.num_offset_fcs - 1:
+                    seq.append(nn.ReLU(inplace=True))
+            self.offset_fc = nn.Sequential(*seq)
+            self.offset_fc[-1].weight.data.zero_()
+            self.offset_fc[-1].bias.data.zero_()
+
+    def forward(self, data, rois):
+        assert data.size(1) == self.out_channels
+        if self.no_trans:
+            offset = data.new_empty(0)
+            return deform_roi_pooling(
+                data, rois, offset, self.spatial_scale, self.out_size,
+                self.out_channels, self.no_trans, self.group_size,
+                self.part_size, self.sample_per_part, self.trans_std)
+        else:
+            n = rois.shape[0]
+            offset = data.new_empty(0)
+            x = deform_roi_pooling(data, rois, offset, self.spatial_scale,
+                                   self.out_size, self.out_channels, True,
+                                   self.group_size, self.part_size,
+                                   self.sample_per_part, self.trans_std)
+            offset = self.offset_fc(x.view(n, -1))
+            offset = offset.view(n, 2, self.out_size, self.out_size)
+            return deform_roi_pooling(
+                data, rois, offset, self.spatial_scale, self.out_size,
+                self.out_channels, self.no_trans, self.group_size,
+                self.part_size, self.sample_per_part, self.trans_std)
+
+
+class ModulatedDeformRoIPoolingPack(DeformRoIPooling):
+
+    def __init__(self,
+                 spatial_scale,
+                 out_size,
+                 out_channels,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0,
+                 num_offset_fcs=3,
+                 num_mask_fcs=2,
+                 deform_fc_channels=1024):
+        super(ModulatedDeformRoIPoolingPack, self).__init__(
+            spatial_scale, out_size, out_channels, no_trans, group_size,
+            part_size, sample_per_part, trans_std)
+
+        self.num_offset_fcs = num_offset_fcs
+        self.num_mask_fcs = num_mask_fcs
+        self.deform_fc_channels = deform_fc_channels
+
+        if not no_trans:
+            offset_fc_seq = []
+            ic = self.out_size * self.out_size * self.out_channels
+            for i in range(self.num_offset_fcs):
+                if i < self.num_offset_fcs - 1:
+                    oc = self.deform_fc_channels
+                else:
+                    oc = self.out_size * self.out_size * 2
+                offset_fc_seq.append(nn.Linear(ic, oc))
+                ic = oc
+                if i < self.num_offset_fcs - 1:
+                    offset_fc_seq.append(nn.ReLU(inplace=True))
+            self.offset_fc = nn.Sequential(*offset_fc_seq)
+            self.offset_fc[-1].weight.data.zero_()
+            self.offset_fc[-1].bias.data.zero_()
+
+            mask_fc_seq = []
+            ic = self.out_size * self.out_size * self.out_channels
+            for i in range(self.num_mask_fcs):
+                if i < self.num_mask_fcs - 1:
+                    oc = self.deform_fc_channels
+                else:
+                    oc = self.out_size * self.out_size
+                mask_fc_seq.append(nn.Linear(ic, oc))
+                ic = oc
+                if i < self.num_mask_fcs - 1:
+                    mask_fc_seq.append(nn.ReLU(inplace=True))
+                else:
+                    mask_fc_seq.append(nn.Sigmoid())
+            self.mask_fc = nn.Sequential(*mask_fc_seq)
+            self.mask_fc[-2].weight.data.zero_()
+            self.mask_fc[-2].bias.data.zero_()
+
+    def forward(self, data, rois):
+        assert data.size(1) == self.out_channels
+        if self.no_trans:
+            offset = data.new_empty(0)
+            return deform_roi_pooling(
+                data, rois, offset, self.spatial_scale, self.out_size,
+                self.out_channels, self.no_trans, self.group_size,
+                self.part_size, self.sample_per_part, self.trans_std)
+        else:
+            n = rois.shape[0]
+            offset = data.new_empty(0)
+            x = deform_roi_pooling(data, rois, offset, self.spatial_scale,
+                                   self.out_size, self.out_channels, True,
+                                   self.group_size, self.part_size,
+                                   self.sample_per_part, self.trans_std)
+            offset = self.offset_fc(x.view(n, -1))
+            offset = offset.view(n, 2, self.out_size, self.out_size)
+            mask = self.mask_fc(x.view(n, -1))
+            mask = mask.view(n, 1, self.out_size, self.out_size)
+            return deform_roi_pooling(
+                data, rois, offset, self.spatial_scale, self.out_size,
+                self.out_channels, self.no_trans, self.group_size,
+                self.part_size, self.sample_per_part, self.trans_std) * mask

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/setup.py

@@ -0,0 +1,19 @@
+import os
+PATH ="{}:{}".format(os.environ['PATH'], "/opt/cuda/bin")
+# os.environ['CUDA_VISIBLE_DEVICES'] = "1"
+os.environ['PATH'] = PATH
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='deform_conv',
+    ext_modules=[
+        CUDAExtension('deform_conv_cuda', [
+            'src/deform_conv_cuda.cpp',
+            'src/deform_conv_cuda_kernel.cu',
+        ]),
+        CUDAExtension('deform_pool_cuda', [
+            'src/deform_pool_cuda.cpp', 'src/deform_pool_cuda_kernel.cu'
+        ]),
+    ],
+    cmdclass={'build_ext': BuildExtension})

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/src/deform_conv_cuda.cpp

@@ -0,0 +1,695 @@
+// modify from
+// https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/deform_conv_cuda.c
+
+#include <torch/extension.h>
+
+#include <cmath>
+#include <vector>
+
+void deformable_im2col(const at::Tensor data_im, const at::Tensor data_offset,
+                       const int channels, const int height, const int width,
+                       const int ksize_h, const int ksize_w, const int pad_h,
+                       const int pad_w, const int stride_h, const int stride_w,
+                       const int dilation_h, const int dilation_w,
+                       const int parallel_imgs, const int deformable_group,
+                       at::Tensor data_col);
+
+void deformable_col2im(const at::Tensor data_col, const at::Tensor data_offset,
+                       const int channels, const int height, const int width,
+                       const int ksize_h, const int ksize_w, const int pad_h,
+                       const int pad_w, const int stride_h, const int stride_w,
+                       const int dilation_h, const int dilation_w,
+                       const int parallel_imgs, const int deformable_group,
+                       at::Tensor grad_im);
+
+void deformable_col2im_coord(
+    const at::Tensor data_col, const at::Tensor data_im,
+    const at::Tensor data_offset, const int channels, const int height,
+    const int width, const int ksize_h, const int ksize_w, const int pad_h,
+    const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w, const int parallel_imgs,
+    const int deformable_group, at::Tensor grad_offset);
+
+void modulated_deformable_im2col_cuda(
+    const at::Tensor data_im, const at::Tensor data_offset,
+    const at::Tensor data_mask, const int batch_size, const int channels,
+    const int height_im, const int width_im, const int height_col,
+    const int width_col, const int kernel_h, const int kenerl_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w, const int deformable_group,
+    at::Tensor data_col);
+
+void modulated_deformable_col2im_cuda(
+    const at::Tensor data_col, const at::Tensor data_offset,
+    const at::Tensor data_mask, const int batch_size, const int channels,
+    const int height_im, const int width_im, const int height_col,
+    const int width_col, const int kernel_h, const int kenerl_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w, const int deformable_group,
+    at::Tensor grad_im);
+
+void modulated_deformable_col2im_coord_cuda(
+    const at::Tensor data_col, const at::Tensor data_im,
+    const at::Tensor data_offset, const at::Tensor data_mask,
+    const int batch_size, const int channels, const int height_im,
+    const int width_im, const int height_col, const int width_col,
+    const int kernel_h, const int kenerl_w, const int pad_h, const int pad_w,
+    const int stride_h, const int stride_w, const int dilation_h,
+    const int dilation_w, const int deformable_group, at::Tensor grad_offset,
+    at::Tensor grad_mask);
+
+void shape_check(at::Tensor input, at::Tensor offset, at::Tensor *gradOutput,
+                 at::Tensor weight, int kH, int kW, int dH, int dW, int padH,
+                 int padW, int dilationH, int dilationW, int group,
+                 int deformable_group) {
+  TORCH_CHECK(weight.ndimension() == 4,
+           "4D weight tensor (nOutputPlane,nInputPlane,kH,kW) expected, "
+           "but got: %s",
+           weight.ndimension());
+
+  TORCH_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
+
+  TORCH_CHECK(kW > 0 && kH > 0,
+           "kernel size should be greater than zero, but got kH: %d kW: %d", kH,
+           kW);
+
+  TORCH_CHECK((weight.size(2) == kH && weight.size(3) == kW),
+           "kernel size should be consistent with weight, ",
+           "but got kH: %d kW: %d weight.size(2): %d, weight.size(3): %d", kH,
+           kW, weight.size(2), weight.size(3));
+
+  TORCH_CHECK(dW > 0 && dH > 0,
+           "stride should be greater than zero, but got dH: %d dW: %d", dH, dW);
+
+  TORCH_CHECK(
+      dilationW > 0 && dilationH > 0,
+      "dilation should be greater than 0, but got dilationH: %d dilationW: %d",
+      dilationH, dilationW);
+
+  int ndim = input.ndimension();
+  int dimf = 0;
+  int dimh = 1;
+  int dimw = 2;
+
+  if (ndim == 4) {
+    dimf++;
+    dimh++;
+    dimw++;
+  }
+
+  TORCH_CHECK(ndim == 3 || ndim == 4, "3D or 4D input tensor expected but got: %s",
+           ndim);
+
+  long nInputPlane = weight.size(1) * group;
+  long inputHeight = input.size(dimh);
+  long inputWidth = input.size(dimw);
+  long nOutputPlane = weight.size(0);
+  long outputHeight =
+      (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+  long outputWidth =
+      (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+
+  TORCH_CHECK(nInputPlane % deformable_group == 0,
+           "input channels must divide deformable group size");
+
+  if (outputWidth < 1 || outputHeight < 1)
+    AT_ERROR(
+        "Given input size: (%ld x %ld x %ld). "
+        "Calculated output size: (%ld x %ld x %ld). Output size is too small",
+        nInputPlane, inputHeight, inputWidth, nOutputPlane, outputHeight,
+        outputWidth);
+
+  TORCH_CHECK(input.size(1) == nInputPlane,
+           "invalid number of input planes, expected: %d, but got: %d",
+           nInputPlane, input.size(1));
+
+  TORCH_CHECK((inputHeight >= kH && inputWidth >= kW),
+           "input image is smaller than kernel");
+
+  TORCH_CHECK((offset.size(2) == outputHeight && offset.size(3) == outputWidth),
+           "invalid spatial size of offset, expected height: %d width: %d, but "
+           "got height: %d width: %d",
+           outputHeight, outputWidth, offset.size(2), offset.size(3));
+
+  TORCH_CHECK((offset.size(1) == deformable_group * 2 * kH * kW),
+           "invalid number of channels of offset");
+
+  if (gradOutput != NULL) {
+    TORCH_CHECK(gradOutput->size(dimf) == nOutputPlane,
+             "invalid number of gradOutput planes, expected: %d, but got: %d",
+             nOutputPlane, gradOutput->size(dimf));
+
+    TORCH_CHECK((gradOutput->size(dimh) == outputHeight &&
+              gradOutput->size(dimw) == outputWidth),
+             "invalid size of gradOutput, expected height: %d width: %d , but "
+             "got height: %d width: %d",
+             outputHeight, outputWidth, gradOutput->size(dimh),
+             gradOutput->size(dimw));
+  }
+}
+
+int deform_conv_forward_cuda(at::Tensor input, at::Tensor weight,
+                             at::Tensor offset, at::Tensor output,
+                             at::Tensor columns, at::Tensor ones, int kW,
+                             int kH, int dW, int dH, int padW, int padH,
+                             int dilationW, int dilationH, int group,
+                             int deformable_group, int im2col_step) {
+  // todo: resize columns to include im2col: done
+  // todo: add im2col_step as input
+  // todo: add new output buffer and transpose it to output (or directly
+  // transpose output) todo: possibly change data indexing because of
+  // parallel_imgs
+
+  shape_check(input, offset, NULL, weight, kH, kW, dH, dW, padH, padW,
+              dilationH, dilationW, group, deformable_group);
+
+  input = input.contiguous();
+  offset = offset.contiguous();
+  weight = weight.contiguous();
+
+  int batch = 1;
+  if (input.ndimension() == 3) {
+    // Force batch
+    batch = 0;
+    input.unsqueeze_(0);
+    offset.unsqueeze_(0);
+  }
+
+  // todo: assert batchsize dividable by im2col_step
+
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
+
+  long nOutputPlane = weight.size(0);
+
+  long outputWidth =
+      (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+  long outputHeight =
+      (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+  TORCH_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
+
+  output = output.view({batchSize / im2col_step, im2col_step, nOutputPlane,
+                        outputHeight, outputWidth});
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
+
+  if (ones.ndimension() != 2 ||
+      ones.size(0) * ones.size(1) < outputHeight * outputWidth) {
+    ones = at::ones({outputHeight, outputWidth}, input.options());
+  }
+
+  input = input.view({batchSize / im2col_step, im2col_step, nInputPlane,
+                      inputHeight, inputWidth});
+  offset =
+      offset.view({batchSize / im2col_step, im2col_step,
+                   deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  at::Tensor output_buffer =
+      at::zeros({batchSize / im2col_step, nOutputPlane,
+                 im2col_step * outputHeight, outputWidth},
+                output.options());
+
+  output_buffer = output_buffer.view(
+      {output_buffer.size(0), group, output_buffer.size(1) / group,
+       output_buffer.size(2), output_buffer.size(3)});
+
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
+    deformable_im2col(input[elt], offset[elt], nInputPlane, inputHeight,
+                      inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
+                      dilationW, im2col_step, deformable_group, columns);
+
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
+
+    for (int g = 0; g < group; g++) {
+      output_buffer[elt][g] = output_buffer[elt][g]
+                                  .flatten(1)
+                                  .addmm_(weight[g].flatten(1), columns[g])
+                                  .view_as(output_buffer[elt][g]);
+    }
+  }
+
+  output_buffer = output_buffer.view(
+      {output_buffer.size(0), output_buffer.size(1) * output_buffer.size(2),
+       output_buffer.size(3), output_buffer.size(4)});
+
+  output_buffer = output_buffer.view({batchSize / im2col_step, nOutputPlane,
+                                      im2col_step, outputHeight, outputWidth});
+  output_buffer.transpose_(1, 2);
+  output.copy_(output_buffer);
+  output = output.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+  input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+  offset = offset.view(
+      {batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  if (batch == 0) {
+    output = output.view({nOutputPlane, outputHeight, outputWidth});
+    input = input.view({nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
+  }
+
+  return 1;
+}
+
+int deform_conv_backward_input_cuda(at::Tensor input, at::Tensor offset,
+                                    at::Tensor gradOutput, at::Tensor gradInput,
+                                    at::Tensor gradOffset, at::Tensor weight,
+                                    at::Tensor columns, int kW, int kH, int dW,
+                                    int dH, int padW, int padH, int dilationW,
+                                    int dilationH, int group,
+                                    int deformable_group, int im2col_step) {
+  shape_check(input, offset, &gradOutput, weight, kH, kW, dH, dW, padH, padW,
+              dilationH, dilationW, group, deformable_group);
+
+  input = input.contiguous();
+  offset = offset.contiguous();
+  gradOutput = gradOutput.contiguous();
+  weight = weight.contiguous();
+
+  int batch = 1;
+
+  if (input.ndimension() == 3) {
+    // Force batch
+    batch = 0;
+    input = input.view({1, input.size(0), input.size(1), input.size(2)});
+    offset = offset.view({1, offset.size(0), offset.size(1), offset.size(2)});
+    gradOutput = gradOutput.view(
+        {1, gradOutput.size(0), gradOutput.size(1), gradOutput.size(2)});
+  }
+
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
+
+  long nOutputPlane = weight.size(0);
+
+  long outputWidth =
+      (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+  long outputHeight =
+      (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+  TORCH_CHECK((offset.size(0) == batchSize), 3, "invalid batch size of offset");
+  gradInput = gradInput.view({batchSize, nInputPlane, inputHeight, inputWidth});
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
+
+  // change order of grad output
+  gradOutput = gradOutput.view({batchSize / im2col_step, im2col_step,
+                                nOutputPlane, outputHeight, outputWidth});
+  gradOutput.transpose_(1, 2);
+
+  gradInput = gradInput.view({batchSize / im2col_step, im2col_step, nInputPlane,
+                              inputHeight, inputWidth});
+  input = input.view({batchSize / im2col_step, im2col_step, nInputPlane,
+                      inputHeight, inputWidth});
+  gradOffset = gradOffset.view({batchSize / im2col_step, im2col_step,
+                                deformable_group * 2 * kH * kW, outputHeight,
+                                outputWidth});
+  offset =
+      offset.view({batchSize / im2col_step, im2col_step,
+                   deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
+    // divide into groups
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
+    gradOutput = gradOutput.view(
+        {gradOutput.size(0), group, gradOutput.size(1) / group,
+         gradOutput.size(2), gradOutput.size(3), gradOutput.size(4)});
+
+    for (int g = 0; g < group; g++) {
+      columns[g] = columns[g].addmm_(weight[g].flatten(1).transpose(0, 1),
+                                     gradOutput[elt][g].flatten(1), 0.0f, 1.0f);
+    }
+
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    gradOutput = gradOutput.view(
+        {gradOutput.size(0), gradOutput.size(1) * gradOutput.size(2),
+         gradOutput.size(3), gradOutput.size(4), gradOutput.size(5)});
+
+    deformable_col2im_coord(columns, input[elt], offset[elt], nInputPlane,
+                            inputHeight, inputWidth, kH, kW, padH, padW, dH, dW,
+                            dilationH, dilationW, im2col_step, deformable_group,
+                            gradOffset[elt]);
+
+    deformable_col2im(columns, offset[elt], nInputPlane, inputHeight,
+                      inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
+                      dilationW, im2col_step, deformable_group, gradInput[elt]);
+  }
+
+  gradOutput.transpose_(1, 2);
+  gradOutput =
+      gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+  gradInput = gradInput.view({batchSize, nInputPlane, inputHeight, inputWidth});
+  input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+  gradOffset = gradOffset.view(
+      {batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+  offset = offset.view(
+      {batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  if (batch == 0) {
+    gradOutput = gradOutput.view({nOutputPlane, outputHeight, outputWidth});
+    input = input.view({nInputPlane, inputHeight, inputWidth});
+    gradInput = gradInput.view({nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
+    gradOffset =
+        gradOffset.view({offset.size(1), offset.size(2), offset.size(3)});
+  }
+
+  return 1;
+}
+
+int deform_conv_backward_parameters_cuda(
+    at::Tensor input, at::Tensor offset, at::Tensor gradOutput,
+    at::Tensor gradWeight,  // at::Tensor gradBias,
+    at::Tensor columns, at::Tensor ones, int kW, int kH, int dW, int dH,
+    int padW, int padH, int dilationW, int dilationH, int group,
+    int deformable_group, float scale, int im2col_step) {
+  // todo: transpose and reshape outGrad
+  // todo: reshape columns
+  // todo: add im2col_step as input
+
+  shape_check(input, offset, &gradOutput, gradWeight, kH, kW, dH, dW, padH,
+              padW, dilationH, dilationW, group, deformable_group);
+
+  input = input.contiguous();
+  offset = offset.contiguous();
+  gradOutput = gradOutput.contiguous();
+
+  int batch = 1;
+
+  if (input.ndimension() == 3) {
+    // Force batch
+    batch = 0;
+    input = input.view(
+        at::IntList({1, input.size(0), input.size(1), input.size(2)}));
+    gradOutput = gradOutput.view(
+        {1, gradOutput.size(0), gradOutput.size(1), gradOutput.size(2)});
+  }
+
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
+
+  long nOutputPlane = gradWeight.size(0);
+
+  long outputWidth =
+      (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+  long outputHeight =
+      (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+  TORCH_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
+
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
+
+  gradOutput = gradOutput.view({batchSize / im2col_step, im2col_step,
+                                nOutputPlane, outputHeight, outputWidth});
+  gradOutput.transpose_(1, 2);
+
+  at::Tensor gradOutputBuffer = at::zeros_like(gradOutput);
+  gradOutputBuffer =
+      gradOutputBuffer.view({batchSize / im2col_step, nOutputPlane, im2col_step,
+                             outputHeight, outputWidth});
+  gradOutputBuffer.copy_(gradOutput);
+  gradOutputBuffer =
+      gradOutputBuffer.view({batchSize / im2col_step, nOutputPlane,
+                             im2col_step * outputHeight, outputWidth});
+
+  gradOutput.transpose_(1, 2);
+  gradOutput =
+      gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+  input = input.view({batchSize / im2col_step, im2col_step, nInputPlane,
+                      inputHeight, inputWidth});
+  offset =
+      offset.view({batchSize / im2col_step, im2col_step,
+                   deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
+    deformable_im2col(input[elt], offset[elt], nInputPlane, inputHeight,
+                      inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
+                      dilationW, im2col_step, deformable_group, columns);
+
+    // divide into group
+    gradOutputBuffer = gradOutputBuffer.view(
+        {gradOutputBuffer.size(0), group, gradOutputBuffer.size(1) / group,
+         gradOutputBuffer.size(2), gradOutputBuffer.size(3)});
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    gradWeight =
+        gradWeight.view({group, gradWeight.size(0) / group, gradWeight.size(1),
+                         gradWeight.size(2), gradWeight.size(3)});
+
+    for (int g = 0; g < group; g++) {
+      gradWeight[g] = gradWeight[g]
+                          .flatten(1)
+                          .addmm_(gradOutputBuffer[elt][g].flatten(1),
+                                  columns[g].transpose(1, 0), 1.0, scale)
+                          .view_as(gradWeight[g]);
+    }
+    gradOutputBuffer = gradOutputBuffer.view(
+        {gradOutputBuffer.size(0),
+         gradOutputBuffer.size(1) * gradOutputBuffer.size(2),
+         gradOutputBuffer.size(3), gradOutputBuffer.size(4)});
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    gradWeight = gradWeight.view({gradWeight.size(0) * gradWeight.size(1),
+                                  gradWeight.size(2), gradWeight.size(3),
+                                  gradWeight.size(4)});
+  }
+
+  input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+  offset = offset.view(
+      {batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+  if (batch == 0) {
+    gradOutput = gradOutput.view({nOutputPlane, outputHeight, outputWidth});
+    input = input.view({nInputPlane, inputHeight, inputWidth});
+  }
+
+  return 1;
+}
+
+void modulated_deform_conv_cuda_forward(
+    at::Tensor input, at::Tensor weight, at::Tensor bias, at::Tensor ones,
+    at::Tensor offset, at::Tensor mask, at::Tensor output, at::Tensor columns,
+    int kernel_h, int kernel_w, const int stride_h, const int stride_w,
+    const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int deformable_group,
+    const bool with_bias) {
+  TORCH_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+  TORCH_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
+
+  const int batch = input.size(0);
+  const int channels = input.size(1);
+  const int height = input.size(2);
+  const int width = input.size(3);
+
+  const int channels_out = weight.size(0);
+  const int channels_kernel = weight.size(1);
+  const int kernel_h_ = weight.size(2);
+  const int kernel_w_ = weight.size(3);
+
+  if (kernel_h_ != kernel_h || kernel_w_ != kernel_w)
+    AT_ERROR("Input shape and kernel shape wont match: (%d x %d vs %d x %d).",
+             kernel_h_, kernel_w, kernel_h_, kernel_w_);
+  if (channels != channels_kernel * group)
+    AT_ERROR("Input shape and kernel channels wont match: (%d vs %d).",
+             channels, channels_kernel * group);
+
+  const int height_out =
+      (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
+  const int width_out =
+      (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  if (ones.ndimension() != 2 ||
+      ones.size(0) * ones.size(1) < height_out * width_out) {
+    // Resize plane and fill with ones...
+    ones = at::ones({height_out, width_out}, input.options());
+  }
+
+  // resize output
+  output = output.view({batch, channels_out, height_out, width_out}).zero_();
+  // resize temporary columns
+  columns =
+      at::zeros({channels * kernel_h * kernel_w, 1 * height_out * width_out},
+                input.options());
+
+  output = output.view({output.size(0), group, output.size(1) / group,
+                        output.size(2), output.size(3)});
+
+  for (int b = 0; b < batch; b++) {
+    modulated_deformable_im2col_cuda(
+        input[b], offset[b], mask[b], 1, channels, height, width, height_out,
+        width_out, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
+        dilation_h, dilation_w, deformable_group, columns);
+
+    // divide into group
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+
+    for (int g = 0; g < group; g++) {
+      output[b][g] = output[b][g]
+                         .flatten(1)
+                         .addmm_(weight[g].flatten(1), columns[g])
+                         .view_as(output[b][g]);
+    }
+
+    weight = weight.view({weight.size(0) * weight.size(1), weight.size(2),
+                          weight.size(3), weight.size(4)});
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+  }
+
+  output = output.view({output.size(0), output.size(1) * output.size(2),
+                        output.size(3), output.size(4)});
+
+  if (with_bias) {
+    output += bias.view({1, bias.size(0), 1, 1});
+  }
+}
+
+void modulated_deform_conv_cuda_backward(
+    at::Tensor input, at::Tensor weight, at::Tensor bias, at::Tensor ones,
+    at::Tensor offset, at::Tensor mask, at::Tensor columns,
+    at::Tensor grad_input, at::Tensor grad_weight, at::Tensor grad_bias,
+    at::Tensor grad_offset, at::Tensor grad_mask, at::Tensor grad_output,
+    int kernel_h, int kernel_w, int stride_h, int stride_w, int pad_h,
+    int pad_w, int dilation_h, int dilation_w, int group, int deformable_group,
+    const bool with_bias) {
+  TORCH_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+  TORCH_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
+
+  const int batch = input.size(0);
+  const int channels = input.size(1);
+  const int height = input.size(2);
+  const int width = input.size(3);
+
+  const int channels_kernel = weight.size(1);
+  const int kernel_h_ = weight.size(2);
+  const int kernel_w_ = weight.size(3);
+  if (kernel_h_ != kernel_h || kernel_w_ != kernel_w)
+    AT_ERROR("Input shape and kernel shape wont match: (%d x %d vs %d x %d).",
+             kernel_h_, kernel_w, kernel_h_, kernel_w_);
+  if (channels != channels_kernel * group)
+    AT_ERROR("Input shape and kernel channels wont match: (%d vs %d).",
+             channels, channels_kernel * group);
+
+  const int height_out =
+      (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
+  const int width_out =
+      (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  if (ones.ndimension() != 2 ||
+      ones.size(0) * ones.size(1) < height_out * width_out) {
+    // Resize plane and fill with ones...
+    ones = at::ones({height_out, width_out}, input.options());
+  }
+
+  grad_input = grad_input.view({batch, channels, height, width});
+  columns = at::zeros({channels * kernel_h * kernel_w, height_out * width_out},
+                      input.options());
+
+  grad_output =
+      grad_output.view({grad_output.size(0), group, grad_output.size(1) / group,
+                        grad_output.size(2), grad_output.size(3)});
+
+  for (int b = 0; b < batch; b++) {
+    // divide int group
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
+
+    for (int g = 0; g < group; g++) {
+      columns[g].addmm_(weight[g].flatten(1).transpose(0, 1),
+                        grad_output[b][g].flatten(1), 0.0f, 1.0f);
+    }
+
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    weight = weight.view({weight.size(0) * weight.size(1), weight.size(2),
+                          weight.size(3), weight.size(4)});
+
+    // gradient w.r.t. input coordinate data
+    modulated_deformable_col2im_coord_cuda(
+        columns, input[b], offset[b], mask[b], 1, channels, height, width,
+        height_out, width_out, kernel_h, kernel_w, pad_h, pad_w, stride_h,
+        stride_w, dilation_h, dilation_w, deformable_group, grad_offset[b],
+        grad_mask[b]);
+    // gradient w.r.t. input data
+    modulated_deformable_col2im_cuda(
+        columns, offset[b], mask[b], 1, channels, height, width, height_out,
+        width_out, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
+        dilation_h, dilation_w, deformable_group, grad_input[b]);
+
+    // gradient w.r.t. weight, dWeight should accumulate across the batch and
+    // group
+    modulated_deformable_im2col_cuda(
+        input[b], offset[b], mask[b], 1, channels, height, width, height_out,
+        width_out, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
+        dilation_h, dilation_w, deformable_group, columns);
+
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    grad_weight = grad_weight.view({group, grad_weight.size(0) / group,
+                                    grad_weight.size(1), grad_weight.size(2),
+                                    grad_weight.size(3)});
+    if (with_bias)
+      grad_bias = grad_bias.view({group, grad_bias.size(0) / group});
+
+    for (int g = 0; g < group; g++) {
+      grad_weight[g] =
+          grad_weight[g]
+              .flatten(1)
+              .addmm_(grad_output[b][g].flatten(1), columns[g].transpose(0, 1))
+              .view_as(grad_weight[g]);
+      if (with_bias) {
+        grad_bias[g] =
+            grad_bias[g]
+                .view({-1, 1})
+                .addmm_(grad_output[b][g].flatten(1), ones.view({-1, 1}))
+                .view(-1);
+      }
+    }
+
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    grad_weight = grad_weight.view({grad_weight.size(0) * grad_weight.size(1),
+                                    grad_weight.size(2), grad_weight.size(3),
+                                    grad_weight.size(4)});
+    if (with_bias)
+      grad_bias = grad_bias.view({grad_bias.size(0) * grad_bias.size(1)});
+  }
+  grad_output = grad_output.view({grad_output.size(0) * grad_output.size(1),
+                                  grad_output.size(2), grad_output.size(3),
+                                  grad_output.size(4)});
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("deform_conv_forward_cuda", &deform_conv_forward_cuda,
+        "deform forward (CUDA)");
+  m.def("deform_conv_backward_input_cuda", &deform_conv_backward_input_cuda,
+        "deform_conv_backward_input (CUDA)");
+  m.def("deform_conv_backward_parameters_cuda",
+        &deform_conv_backward_parameters_cuda,
+        "deform_conv_backward_parameters (CUDA)");
+  m.def("modulated_deform_conv_cuda_forward",
+        &modulated_deform_conv_cuda_forward,
+        "modulated deform conv forward (CUDA)");
+  m.def("modulated_deform_conv_cuda_backward",
+        &modulated_deform_conv_cuda_backward,
+        "modulated deform conv backward (CUDA)");
+}

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/src/deform_conv_cuda_kernel.cu

@@ -0,0 +1,866 @@
+/*!
+ ******************* BEGIN Caffe Copyright Notice and Disclaimer ****************
+ *
+ * COPYRIGHT
+ *
+ * All contributions by the University of California:
+ * Copyright (c) 2014-2017 The Regents of the University of California (Regents)
+ * All rights reserved.
+ *
+ * All other contributions:
+ * Copyright (c) 2014-2017, the respective contributors
+ * All rights reserved.
+ *
+ * Caffe uses a shared copyright model: each contributor holds copyright over
+ * their contributions to Caffe. The project versioning records all such
+ * contribution and copyright details. If a contributor wants to further mark
+ * their specific copyright on a particular contribution, they should indicate
+ * their copyright solely in the commit message of the change when it is
+ * committed.
+ *
+ * LICENSE
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * CONTRIBUTION AGREEMENT
+ *
+ * By contributing to the BVLC/caffe repository through pull-request, comment,
+ * or otherwise, the contributor releases their content to the
+ * license and copyright terms herein.
+ *
+ ***************** END Caffe Copyright Notice and Disclaimer ********************
+ *
+ * Copyright (c) 2018 Microsoft
+ * Licensed under The MIT License [see LICENSE for details]
+ * \file modulated_deformable_im2col.cuh
+ * \brief Function definitions of converting an image to
+ * column matrix based on kernel, padding, dilation, and offset.
+ * These functions are mainly used in deformable convolution operators.
+ * \ref: https://arxiv.org/abs/1703.06211
+ * \author Yuwen Xiong, Haozhi Qi, Jifeng Dai, Xizhou Zhu, Han Hu, Dazhi Cheng
+ */
+
+// modify from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/deform_conv_cuda_kernel.cu
+
+#include <ATen/ATen.h>
+#include <THC/THCAtomics.cuh>
+#include <stdio.h>
+#include <math.h>
+#include <float.h>
+
+using namespace at;
+
+#define CUDA_KERNEL_LOOP(i, n)                                 \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
+       i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+const int kMaxGridNum = 65535;
+
+inline int GET_BLOCKS(const int N)
+{
+  return std::min(kMaxGridNum, (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS);
+}
+
+template <typename scalar_t>
+__device__ scalar_t deformable_im2col_bilinear(const scalar_t *bottom_data, const int data_width,
+                                               const int height, const int width, scalar_t h, scalar_t w)
+{
+
+  int h_low = floor(h);
+  int w_low = floor(w);
+  int h_high = h_low + 1;
+  int w_high = w_low + 1;
+
+  scalar_t lh = h - h_low;
+  scalar_t lw = w - w_low;
+  scalar_t hh = 1 - lh, hw = 1 - lw;
+
+  scalar_t v1 = 0;
+  if (h_low >= 0 && w_low >= 0)
+    v1 = bottom_data[h_low * data_width + w_low];
+  scalar_t v2 = 0;
+  if (h_low >= 0 && w_high <= width - 1)
+    v2 = bottom_data[h_low * data_width + w_high];
+  scalar_t v3 = 0;
+  if (h_high <= height - 1 && w_low >= 0)
+    v3 = bottom_data[h_high * data_width + w_low];
+  scalar_t v4 = 0;
+  if (h_high <= height - 1 && w_high <= width - 1)
+    v4 = bottom_data[h_high * data_width + w_high];
+
+  scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  return val;
+}
+
+template <typename scalar_t>
+__device__ scalar_t get_gradient_weight(scalar_t argmax_h, scalar_t argmax_w,
+                                        const int h, const int w, const int height, const int width)
+{
+
+  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
+  {
+    //empty
+    return 0;
+  }
+
+  int argmax_h_low = floor(argmax_h);
+  int argmax_w_low = floor(argmax_w);
+  int argmax_h_high = argmax_h_low + 1;
+  int argmax_w_high = argmax_w_low + 1;
+
+  scalar_t weight = 0;
+  if (h == argmax_h_low && w == argmax_w_low)
+    weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);
+  if (h == argmax_h_low && w == argmax_w_high)
+    weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);
+  if (h == argmax_h_high && w == argmax_w_low)
+    weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);
+  if (h == argmax_h_high && w == argmax_w_high)
+    weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);
+  return weight;
+}
+
+template <typename scalar_t>
+__device__ scalar_t get_coordinate_weight(scalar_t argmax_h, scalar_t argmax_w,
+                                          const int height, const int width, const scalar_t *im_data,
+                                          const int data_width, const int bp_dir)
+{
+
+  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
+  {
+    //empty
+    return 0;
+  }
+
+  int argmax_h_low = floor(argmax_h);
+  int argmax_w_low = floor(argmax_w);
+  int argmax_h_high = argmax_h_low + 1;
+  int argmax_w_high = argmax_w_low + 1;
+
+  scalar_t weight = 0;
+
+  if (bp_dir == 0)
+  {
+    if (argmax_h_low >= 0 && argmax_w_low >= 0)
+      weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];
+    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
+      weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];
+    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
+      weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];
+    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
+      weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];
+  }
+  else if (bp_dir == 1)
+  {
+    if (argmax_h_low >= 0 && argmax_w_low >= 0)
+      weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];
+    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
+      weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];
+    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
+      weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];
+    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
+      weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];
+  }
+
+  return weight;
+}
+
+template <typename scalar_t>
+__global__ void deformable_im2col_gpu_kernel(const int n, const scalar_t *data_im, const scalar_t *data_offset,
+                                             const int height, const int width, const int kernel_h, const int kernel_w,
+                                             const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+                                             const int dilation_h, const int dilation_w, const int channel_per_deformable_group,
+                                             const int batch_size, const int num_channels, const int deformable_group,
+                                             const int height_col, const int width_col,
+                                             scalar_t *data_col)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    // index index of output matrix
+    const int w_col = index % width_col;
+    const int h_col = (index / width_col) % height_col;
+    const int b_col = (index / width_col / height_col) % batch_size;
+    const int c_im = (index / width_col / height_col) / batch_size;
+    const int c_col = c_im * kernel_h * kernel_w;
+
+    // compute deformable group index
+    const int deformable_group_index = c_im / channel_per_deformable_group;
+
+    const int h_in = h_col * stride_h - pad_h;
+    const int w_in = w_col * stride_w - pad_w;
+    scalar_t *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
+    //const scalar_t* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
+    const scalar_t *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
+    const scalar_t *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
+
+    for (int i = 0; i < kernel_h; ++i)
+    {
+      for (int j = 0; j < kernel_w; ++j)
+      {
+        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
+        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
+        const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+        const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+        scalar_t val = static_cast<scalar_t>(0);
+        const scalar_t h_im = h_in + i * dilation_h + offset_h;
+        const scalar_t w_im = w_in + j * dilation_w + offset_w;
+        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
+        {
+          //const scalar_t map_h = i * dilation_h + offset_h;
+          //const scalar_t map_w = j * dilation_w + offset_w;
+          //const int cur_height = height - h_in;
+          //const int cur_width = width - w_in;
+          //val = deformable_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
+          val = deformable_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
+        }
+        *data_col_ptr = val;
+        data_col_ptr += batch_size * height_col * width_col;
+      }
+    }
+  }
+}
+
+void deformable_im2col(
+    const at::Tensor data_im, const at::Tensor data_offset, const int channels,
+    const int height, const int width, const int ksize_h, const int ksize_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w, const int parallel_imgs,
+    const int deformable_group, at::Tensor data_col)
+{
+  // num_axes should be smaller than block size
+  // todo: check parallel_imgs is correctly passed in
+  int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
+  int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
+  int num_kernels = channels * height_col * width_col * parallel_imgs;
+  int channel_per_deformable_group = channels / deformable_group;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_im.type(), "deformable_im2col_gpu", ([&] {
+        const scalar_t *data_im_ = data_im.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        scalar_t *data_col_ = data_col.data<scalar_t>();
+
+        deformable_im2col_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_im_, data_offset_, height, width, ksize_h, ksize_w,
+            pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
+            channel_per_deformable_group, parallel_imgs, channels, deformable_group,
+            height_col, width_col, data_col_);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in deformable_im2col: %s\n", cudaGetErrorString(err));
+  }
+}
+
+template <typename scalar_t>
+__global__ void deformable_col2im_gpu_kernel(
+    const int n, const scalar_t *data_col, const scalar_t *data_offset,
+    const int channels, const int height, const int width,
+    const int kernel_h, const int kernel_w,
+    const int pad_h, const int pad_w,
+    const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w,
+    const int channel_per_deformable_group,
+    const int batch_size, const int deformable_group,
+    const int height_col, const int width_col,
+    scalar_t *grad_im)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    const int j = (index / width_col / height_col / batch_size) % kernel_w;
+    const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;
+    const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;
+    // compute the start and end of the output
+
+    const int deformable_group_index = c / channel_per_deformable_group;
+
+    int w_out = index % width_col;
+    int h_out = (index / width_col) % height_col;
+    int b = (index / width_col / height_col) % batch_size;
+    int w_in = w_out * stride_w - pad_w;
+    int h_in = h_out * stride_h - pad_h;
+
+    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) *
+                                                        2 * kernel_h * kernel_w * height_col * width_col;
+    const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;
+    const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;
+    const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+    const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+    const scalar_t cur_inv_h_data = h_in + i * dilation_h + offset_h;
+    const scalar_t cur_inv_w_data = w_in + j * dilation_w + offset_w;
+
+    const scalar_t cur_top_grad = data_col[index];
+    const int cur_h = (int)cur_inv_h_data;
+    const int cur_w = (int)cur_inv_w_data;
+    for (int dy = -2; dy <= 2; dy++)
+    {
+      for (int dx = -2; dx <= 2; dx++)
+      {
+        if (cur_h + dy >= 0 && cur_h + dy < height &&
+            cur_w + dx >= 0 && cur_w + dx < width &&
+            abs(cur_inv_h_data - (cur_h + dy)) < 1 &&
+            abs(cur_inv_w_data - (cur_w + dx)) < 1)
+        {
+          int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;
+          scalar_t weight = get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);
+          atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);
+        }
+      }
+    }
+  }
+}
+
+void deformable_col2im(
+    const at::Tensor data_col, const at::Tensor data_offset, const int channels,
+    const int height, const int width, const int ksize_h,
+    const int ksize_w, const int pad_h, const int pad_w,
+    const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w,
+    const int parallel_imgs, const int deformable_group,
+    at::Tensor grad_im)
+{
+
+  // todo: make sure parallel_imgs is passed in correctly
+  int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
+  int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
+  int num_kernels = channels * ksize_h * ksize_w * height_col * width_col * parallel_imgs;
+  int channel_per_deformable_group = channels / deformable_group;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.type(), "deformable_col2im_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        scalar_t *grad_im_ = grad_im.data<scalar_t>();
+
+        deformable_col2im_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_col_, data_offset_, channels, height, width, ksize_h,
+            ksize_w, pad_h, pad_w, stride_h, stride_w,
+            dilation_h, dilation_w, channel_per_deformable_group,
+            parallel_imgs, deformable_group, height_col, width_col, grad_im_);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in deformable_col2im: %s\n", cudaGetErrorString(err));
+  }
+}
+
+template <typename scalar_t>
+__global__ void deformable_col2im_coord_gpu_kernel(const int n, const scalar_t *data_col,
+                                                   const scalar_t *data_im, const scalar_t *data_offset,
+                                                   const int channels, const int height, const int width,
+                                                   const int kernel_h, const int kernel_w,
+                                                   const int pad_h, const int pad_w,
+                                                   const int stride_h, const int stride_w,
+                                                   const int dilation_h, const int dilation_w,
+                                                   const int channel_per_deformable_group,
+                                                   const int batch_size, const int offset_channels, const int deformable_group,
+                                                   const int height_col, const int width_col, scalar_t *grad_offset)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    scalar_t val = 0;
+    int w = index % width_col;
+    int h = (index / width_col) % height_col;
+    int c = (index / width_col / height_col) % offset_channels;
+    int b = (index / width_col / height_col) / offset_channels;
+    // compute the start and end of the output
+
+    const int deformable_group_index = c / (2 * kernel_h * kernel_w);
+    const int col_step = kernel_h * kernel_w;
+    int cnt = 0;
+    const scalar_t *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group *
+                                                  batch_size * width_col * height_col;
+    const scalar_t *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) *
+                                                channel_per_deformable_group / kernel_h / kernel_w * height * width;
+    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 *
+                                                        kernel_h * kernel_w * height_col * width_col;
+
+    const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;
+
+    for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)
+    {
+      const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;
+      const int bp_dir = offset_c % 2;
+
+      int j = (col_pos / width_col / height_col / batch_size) % kernel_w;
+      int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;
+      int w_out = col_pos % width_col;
+      int h_out = (col_pos / width_col) % height_col;
+      int w_in = w_out * stride_w - pad_w;
+      int h_in = h_out * stride_h - pad_h;
+      const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);
+      const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);
+      const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+      const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+      scalar_t inv_h = h_in + i * dilation_h + offset_h;
+      scalar_t inv_w = w_in + j * dilation_w + offset_w;
+      if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)
+      {
+        inv_h = inv_w = -2;
+      }
+      const scalar_t weight = get_coordinate_weight(
+          inv_h, inv_w,
+          height, width, data_im_ptr + cnt * height * width, width, bp_dir);
+      val += weight * data_col_ptr[col_pos];
+      cnt += 1;
+    }
+
+    grad_offset[index] = val;
+  }
+}
+
+void deformable_col2im_coord(
+    const at::Tensor data_col, const at::Tensor data_im, const at::Tensor data_offset,
+    const int channels, const int height, const int width, const int ksize_h,
+    const int ksize_w, const int pad_h, const int pad_w, const int stride_h,
+    const int stride_w, const int dilation_h, const int dilation_w,
+    const int parallel_imgs, const int deformable_group, at::Tensor grad_offset)
+{
+
+  int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
+  int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
+  int num_kernels = height_col * width_col * 2 * ksize_h * ksize_w * deformable_group * parallel_imgs;
+  int channel_per_deformable_group = channels * ksize_h * ksize_w / deformable_group;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.type(), "deformable_col2im_coord_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data<scalar_t>();
+        const scalar_t *data_im_ = data_im.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        scalar_t *grad_offset_ = grad_offset.data<scalar_t>();
+
+        deformable_col2im_coord_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_col_, data_im_, data_offset_, channels, height, width,
+            ksize_h, ksize_w, pad_h, pad_w, stride_h, stride_w,
+            dilation_h, dilation_w, channel_per_deformable_group,
+            parallel_imgs, 2 * ksize_h * ksize_w * deformable_group, deformable_group,
+            height_col, width_col, grad_offset_);
+      }));
+}
+
+template <typename scalar_t>
+__device__ scalar_t dmcn_im2col_bilinear(const scalar_t *bottom_data, const int data_width,
+                                         const int height, const int width, scalar_t h, scalar_t w)
+{
+  int h_low = floor(h);
+  int w_low = floor(w);
+  int h_high = h_low + 1;
+  int w_high = w_low + 1;
+
+  scalar_t lh = h - h_low;
+  scalar_t lw = w - w_low;
+  scalar_t hh = 1 - lh, hw = 1 - lw;
+
+  scalar_t v1 = 0;
+  if (h_low >= 0 && w_low >= 0)
+    v1 = bottom_data[h_low * data_width + w_low];
+  scalar_t v2 = 0;
+  if (h_low >= 0 && w_high <= width - 1)
+    v2 = bottom_data[h_low * data_width + w_high];
+  scalar_t v3 = 0;
+  if (h_high <= height - 1 && w_low >= 0)
+    v3 = bottom_data[h_high * data_width + w_low];
+  scalar_t v4 = 0;
+  if (h_high <= height - 1 && w_high <= width - 1)
+    v4 = bottom_data[h_high * data_width + w_high];
+
+  scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  return val;
+}
+
+template <typename scalar_t>
+__device__ scalar_t dmcn_get_gradient_weight(scalar_t argmax_h, scalar_t argmax_w,
+                                             const int h, const int w, const int height, const int width)
+{
+  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
+  {
+    //empty
+    return 0;
+  }
+
+  int argmax_h_low = floor(argmax_h);
+  int argmax_w_low = floor(argmax_w);
+  int argmax_h_high = argmax_h_low + 1;
+  int argmax_w_high = argmax_w_low + 1;
+
+  scalar_t weight = 0;
+  if (h == argmax_h_low && w == argmax_w_low)
+    weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);
+  if (h == argmax_h_low && w == argmax_w_high)
+    weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);
+  if (h == argmax_h_high && w == argmax_w_low)
+    weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);
+  if (h == argmax_h_high && w == argmax_w_high)
+    weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);
+  return weight;
+}
+
+template <typename scalar_t>
+__device__ scalar_t dmcn_get_coordinate_weight(scalar_t argmax_h, scalar_t argmax_w,
+                                               const int height, const int width, const scalar_t *im_data,
+                                               const int data_width, const int bp_dir)
+{
+  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
+  {
+    //empty
+    return 0;
+  }
+
+  int argmax_h_low = floor(argmax_h);
+  int argmax_w_low = floor(argmax_w);
+  int argmax_h_high = argmax_h_low + 1;
+  int argmax_w_high = argmax_w_low + 1;
+
+  scalar_t weight = 0;
+
+  if (bp_dir == 0)
+  {
+    if (argmax_h_low >= 0 && argmax_w_low >= 0)
+      weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];
+    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
+      weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];
+    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
+      weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];
+    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
+      weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];
+  }
+  else if (bp_dir == 1)
+  {
+    if (argmax_h_low >= 0 && argmax_w_low >= 0)
+      weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];
+    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
+      weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];
+    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
+      weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];
+    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
+      weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];
+  }
+
+  return weight;
+}
+
+template <typename scalar_t>
+__global__ void modulated_deformable_im2col_gpu_kernel(const int n,
+                                                       const scalar_t *data_im, const scalar_t *data_offset, const scalar_t *data_mask,
+                                                       const int height, const int width, const int kernel_h, const int kernel_w,
+                                                       const int pad_h, const int pad_w,
+                                                       const int stride_h, const int stride_w,
+                                                       const int dilation_h, const int dilation_w,
+                                                       const int channel_per_deformable_group,
+                                                       const int batch_size, const int num_channels, const int deformable_group,
+                                                       const int height_col, const int width_col,
+                                                       scalar_t *data_col)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    // index index of output matrix
+    const int w_col = index % width_col;
+    const int h_col = (index / width_col) % height_col;
+    const int b_col = (index / width_col / height_col) % batch_size;
+    const int c_im = (index / width_col / height_col) / batch_size;
+    const int c_col = c_im * kernel_h * kernel_w;
+
+    // compute deformable group index
+    const int deformable_group_index = c_im / channel_per_deformable_group;
+
+    const int h_in = h_col * stride_h - pad_h;
+    const int w_in = w_col * stride_w - pad_w;
+
+    scalar_t *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
+    //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
+    const scalar_t *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
+    const scalar_t *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
+
+    const scalar_t *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
+
+    for (int i = 0; i < kernel_h; ++i)
+    {
+      for (int j = 0; j < kernel_w; ++j)
+      {
+        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
+        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
+        const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
+        const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+        const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+        const scalar_t mask = data_mask_ptr[data_mask_hw_ptr];
+        scalar_t val = static_cast<scalar_t>(0);
+        const scalar_t h_im = h_in + i * dilation_h + offset_h;
+        const scalar_t w_im = w_in + j * dilation_w + offset_w;
+        //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
+        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
+        {
+          //const float map_h = i * dilation_h + offset_h;
+          //const float map_w = j * dilation_w + offset_w;
+          //const int cur_height = height - h_in;
+          //const int cur_width = width - w_in;
+          //val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
+          val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
+        }
+        *data_col_ptr = val * mask;
+        data_col_ptr += batch_size * height_col * width_col;
+        //data_col_ptr += height_col * width_col;
+      }
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void modulated_deformable_col2im_gpu_kernel(const int n,
+                                                       const scalar_t *data_col, const scalar_t *data_offset, const scalar_t *data_mask,
+                                                       const int channels, const int height, const int width,
+                                                       const int kernel_h, const int kernel_w,
+                                                       const int pad_h, const int pad_w,
+                                                       const int stride_h, const int stride_w,
+                                                       const int dilation_h, const int dilation_w,
+                                                       const int channel_per_deformable_group,
+                                                       const int batch_size, const int deformable_group,
+                                                       const int height_col, const int width_col,
+                                                       scalar_t *grad_im)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    const int j = (index / width_col / height_col / batch_size) % kernel_w;
+    const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;
+    const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;
+    // compute the start and end of the output
+
+    const int deformable_group_index = c / channel_per_deformable_group;
+
+    int w_out = index % width_col;
+    int h_out = (index / width_col) % height_col;
+    int b = (index / width_col / height_col) % batch_size;
+    int w_in = w_out * stride_w - pad_w;
+    int h_in = h_out * stride_h - pad_h;
+
+    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
+    const scalar_t *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
+    const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;
+    const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;
+    const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out;
+    const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+    const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+    const scalar_t mask = data_mask_ptr[data_mask_hw_ptr];
+    const scalar_t cur_inv_h_data = h_in + i * dilation_h + offset_h;
+    const scalar_t cur_inv_w_data = w_in + j * dilation_w + offset_w;
+
+    const scalar_t cur_top_grad = data_col[index] * mask;
+    const int cur_h = (int)cur_inv_h_data;
+    const int cur_w = (int)cur_inv_w_data;
+    for (int dy = -2; dy <= 2; dy++)
+    {
+      for (int dx = -2; dx <= 2; dx++)
+      {
+        if (cur_h + dy >= 0 && cur_h + dy < height &&
+            cur_w + dx >= 0 && cur_w + dx < width &&
+            abs(cur_inv_h_data - (cur_h + dy)) < 1 &&
+            abs(cur_inv_w_data - (cur_w + dx)) < 1)
+        {
+          int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;
+          scalar_t weight = dmcn_get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);
+          atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);
+        }
+      }
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void modulated_deformable_col2im_coord_gpu_kernel(const int n,
+                                                             const scalar_t *data_col, const scalar_t *data_im,
+                                                             const scalar_t *data_offset, const scalar_t *data_mask,
+                                                             const int channels, const int height, const int width,
+                                                             const int kernel_h, const int kernel_w,
+                                                             const int pad_h, const int pad_w,
+                                                             const int stride_h, const int stride_w,
+                                                             const int dilation_h, const int dilation_w,
+                                                             const int channel_per_deformable_group,
+                                                             const int batch_size, const int offset_channels, const int deformable_group,
+                                                             const int height_col, const int width_col,
+                                                             scalar_t *grad_offset, scalar_t *grad_mask)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    scalar_t val = 0, mval = 0;
+    int w = index % width_col;
+    int h = (index / width_col) % height_col;
+    int c = (index / width_col / height_col) % offset_channels;
+    int b = (index / width_col / height_col) / offset_channels;
+    // compute the start and end of the output
+
+    const int deformable_group_index = c / (2 * kernel_h * kernel_w);
+    const int col_step = kernel_h * kernel_w;
+    int cnt = 0;
+    const scalar_t *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col;
+    const scalar_t *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width;
+    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
+    const scalar_t *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
+
+    const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;
+
+    for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)
+    {
+      const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;
+      const int bp_dir = offset_c % 2;
+
+      int j = (col_pos / width_col / height_col / batch_size) % kernel_w;
+      int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;
+      int w_out = col_pos % width_col;
+      int h_out = (col_pos / width_col) % height_col;
+      int w_in = w_out * stride_w - pad_w;
+      int h_in = h_out * stride_h - pad_h;
+      const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);
+      const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);
+      const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out);
+      const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
+      const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
+      const scalar_t mask = data_mask_ptr[data_mask_hw_ptr];
+      scalar_t inv_h = h_in + i * dilation_h + offset_h;
+      scalar_t inv_w = w_in + j * dilation_w + offset_w;
+      if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)
+      {
+        inv_h = inv_w = -2;
+      }
+      else
+      {
+        mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w);
+      }
+      const scalar_t weight = dmcn_get_coordinate_weight(
+          inv_h, inv_w,
+          height, width, data_im_ptr + cnt * height * width, width, bp_dir);
+      val += weight * data_col_ptr[col_pos] * mask;
+      cnt += 1;
+    }
+    // KERNEL_ASSIGN(grad_offset[index], offset_req, val);
+    grad_offset[index] = val;
+    if (offset_c % 2 == 0)
+      // KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval);
+      grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval;
+  }
+}
+
+void modulated_deformable_im2col_cuda(
+    const at::Tensor data_im, const at::Tensor data_offset, const at::Tensor data_mask,
+    const int batch_size, const int channels, const int height_im, const int width_im,
+    const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w,
+    const int deformable_group, at::Tensor data_col)
+{
+  // num_axes should be smaller than block size
+  const int channel_per_deformable_group = channels / deformable_group;
+  const int num_kernels = channels * batch_size * height_col * width_col;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_im.type(), "modulated_deformable_im2col_gpu", ([&] {
+        const scalar_t *data_im_ = data_im.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        const scalar_t *data_mask_ = data_mask.data<scalar_t>();
+        scalar_t *data_col_ = data_col.data<scalar_t>();
+
+        modulated_deformable_im2col_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_im_, data_offset_, data_mask_, height_im, width_im, kernel_h, kenerl_w,
+            pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,
+            batch_size, channels, deformable_group, height_col, width_col, data_col_);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    // printf("error in modulated_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+  }
+}
+
+void modulated_deformable_col2im_cuda(
+    const at::Tensor data_col, const at::Tensor data_offset, const at::Tensor data_mask,
+    const int batch_size, const int channels, const int height_im, const int width_im,
+    const int height_col, const int width_col, const int kernel_h, const int kernel_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w,
+    const int deformable_group, at::Tensor grad_im)
+{
+
+  const int channel_per_deformable_group = channels / deformable_group;
+  const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.type(), "modulated_deformable_col2im_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        const scalar_t *data_mask_ = data_mask.data<scalar_t>();
+        scalar_t *grad_im_ = grad_im.data<scalar_t>();
+
+        modulated_deformable_col2im_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_col_, data_offset_, data_mask_, channels, height_im, width_im,
+            kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w,
+            dilation_h, dilation_w, channel_per_deformable_group,
+            batch_size, deformable_group, height_col, width_col, grad_im_);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in modulated_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
+  }
+}
+
+void modulated_deformable_col2im_coord_cuda(
+    const at::Tensor data_col, const at::Tensor data_im, const at::Tensor data_offset, const at::Tensor data_mask,
+    const int batch_size, const int channels, const int height_im, const int width_im,
+    const int height_col, const int width_col, const int kernel_h, const int kernel_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w,
+    const int deformable_group,
+    at::Tensor grad_offset, at::Tensor grad_mask)
+{
+  const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group;
+  const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.type(), "modulated_deformable_col2im_coord_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data<scalar_t>();
+        const scalar_t *data_im_ = data_im.data<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data<scalar_t>();
+        const scalar_t *data_mask_ = data_mask.data<scalar_t>();
+        scalar_t *grad_offset_ = grad_offset.data<scalar_t>();
+        scalar_t *grad_mask_ = grad_mask.data<scalar_t>();
+
+        modulated_deformable_col2im_coord_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
+            num_kernels, data_col_, data_im_, data_offset_, data_mask_, channels, height_im, width_im,
+            kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
+            dilation_h, dilation_w, channel_per_deformable_group,
+            batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col,
+            grad_offset_, grad_mask_);
+      }));
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in modulated_deformable_col2im_coord_cuda: %s\n", cudaGetErrorString(err));
+  }
+}

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/src/deform_pool_cuda.cpp

@@ -0,0 +1,87 @@
+// modify from
+// https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/modulated_dcn_cuda.c
+
+// based on
+// author: Charles Shang
+// https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu
+
+#include <torch/extension.h>
+
+#include <cmath>
+#include <vector>
+
+void DeformablePSROIPoolForward(
+    const at::Tensor data, const at::Tensor bbox, const at::Tensor trans,
+    at::Tensor out, at::Tensor top_count, const int batch, const int channels,
+    const int height, const int width, const int num_bbox,
+    const int channels_trans, const int no_trans, const float spatial_scale,
+    const int output_dim, const int group_size, const int pooled_size,
+    const int part_size, const int sample_per_part, const float trans_std);
+
+void DeformablePSROIPoolBackwardAcc(
+    const at::Tensor out_grad, const at::Tensor data, const at::Tensor bbox,
+    const at::Tensor trans, const at::Tensor top_count, at::Tensor in_grad,
+    at::Tensor trans_grad, const int batch, const int channels,
+    const int height, const int width, const int num_bbox,
+    const int channels_trans, const int no_trans, const float spatial_scale,
+    const int output_dim, const int group_size, const int pooled_size,
+    const int part_size, const int sample_per_part, const float trans_std);
+
+void deform_psroi_pooling_cuda_forward(
+    at::Tensor input, at::Tensor bbox, at::Tensor trans, at::Tensor out,
+    at::Tensor top_count, const int no_trans, const float spatial_scale,
+    const int output_dim, const int group_size, const int pooled_size,
+    const int part_size, const int sample_per_part, const float trans_std) {
+  TORCH_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+
+  const int batch = input.size(0);
+  const int channels = input.size(1);
+  const int height = input.size(2);
+  const int width = input.size(3);
+  const int channels_trans = no_trans ? 2 : trans.size(1);
+
+  const int num_bbox = bbox.size(0);
+  if (num_bbox != out.size(0))
+    AT_ERROR("Output shape and bbox number wont match: (%d vs %d).",
+             out.size(0), num_bbox);
+
+  DeformablePSROIPoolForward(
+      input, bbox, trans, out, top_count, batch, channels, height, width,
+      num_bbox, channels_trans, no_trans, spatial_scale, output_dim, group_size,
+      pooled_size, part_size, sample_per_part, trans_std);
+}
+
+void deform_psroi_pooling_cuda_backward(
+    at::Tensor out_grad, at::Tensor input, at::Tensor bbox, at::Tensor trans,
+    at::Tensor top_count, at::Tensor input_grad, at::Tensor trans_grad,
+    const int no_trans, const float spatial_scale, const int output_dim,
+    const int group_size, const int pooled_size, const int part_size,
+    const int sample_per_part, const float trans_std) {
+  TORCH_CHECK(out_grad.is_contiguous(), "out_grad tensor has to be contiguous");
+  TORCH_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+
+  const int batch = input.size(0);
+  const int channels = input.size(1);
+  const int height = input.size(2);
+  const int width = input.size(3);
+  const int channels_trans = no_trans ? 2 : trans.size(1);
+
+  const int num_bbox = bbox.size(0);
+  if (num_bbox != out_grad.size(0))
+    AT_ERROR("Output shape and bbox number wont match: (%d vs %d).",
+             out_grad.size(0), num_bbox);
+
+  DeformablePSROIPoolBackwardAcc(
+      out_grad, input, bbox, trans, top_count, input_grad, trans_grad, batch,
+      channels, height, width, num_bbox, channels_trans, no_trans,
+      spatial_scale, output_dim, group_size, pooled_size, part_size,
+      sample_per_part, trans_std);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("deform_psroi_pooling_cuda_forward", &deform_psroi_pooling_cuda_forward,
+        "deform psroi pooling forward(CUDA)");
+  m.def("deform_psroi_pooling_cuda_backward",
+        &deform_psroi_pooling_cuda_backward,
+        "deform psroi pooling backward(CUDA)");
+}

IndicPhotoOCR/detection/textbpn/network/backbone/assets/dcn/src/deform_pool_cuda_kernel.cu

@@ -0,0 +1,364 @@
+/*!
+ * Copyright (c) 2017 Microsoft
+ * Licensed under The MIT License [see LICENSE for details]
+ * \file deformable_psroi_pooling.cu
+ * \brief
+ * \author Yi Li, Guodong Zhang, Jifeng Dai
+*/
+/***************** Adapted by Charles Shang *********************/
+// modify from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/cuda/deform_psroi_pooling_cuda.cu
+
+#include <ATen/ATen.h>
+#include <THC/THCAtomics.cuh>
+#include <stdio.h>
+#include <math.h>
+#include <algorithm>
+
+using namespace at;
+
+#define CUDA_KERNEL_LOOP(i, n)                        \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
+       i < (n);                                       \
+       i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N)
+{
+  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
+}
+
+template <typename scalar_t>
+__device__ scalar_t bilinear_interp(
+    const scalar_t *data,
+    const scalar_t x,
+    const scalar_t y,
+    const int width,
+    const int height)
+{
+  int x1 = floor(x);
+  int x2 = ceil(x);
+  int y1 = floor(y);
+  int y2 = ceil(y);
+  scalar_t dist_x = (scalar_t)(x - x1);
+  scalar_t dist_y = (scalar_t)(y - y1);
+  scalar_t value11 = data[y1 * width + x1];
+  scalar_t value12 = data[y2 * width + x1];
+  scalar_t value21 = data[y1 * width + x2];
+  scalar_t value22 = data[y2 * width + x2];
+  scalar_t value = (1 - dist_x) * (1 - dist_y) * value11 + (1 - dist_x) * dist_y * value12 + dist_x * (1 - dist_y) * value21 + dist_x * dist_y * value22;
+  return value;
+}
+
+template <typename scalar_t>
+__global__ void DeformablePSROIPoolForwardKernel(
+    const int count,
+    const scalar_t *bottom_data,
+    const scalar_t spatial_scale,
+    const int channels,
+    const int height, const int width,
+    const int pooled_height, const int pooled_width,
+    const scalar_t *bottom_rois, const scalar_t *bottom_trans,
+    const int no_trans,
+    const scalar_t trans_std,
+    const int sample_per_part,
+    const int output_dim,
+    const int group_size,
+    const int part_size,
+    const int num_classes,
+    const int channels_each_class,
+    scalar_t *top_data,
+    scalar_t *top_count)
+{
+  CUDA_KERNEL_LOOP(index, count)
+  {
+    // The output is in order (n, ctop, ph, pw)
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int ctop = (index / pooled_width / pooled_height) % output_dim;
+    int n = index / pooled_width / pooled_height / output_dim;
+
+    // [start, end) interval for spatial sampling
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 5;
+    int roi_batch_ind = offset_bottom_rois[0];
+    scalar_t roi_start_w = (scalar_t)(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
+    scalar_t roi_start_h = (scalar_t)(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
+    scalar_t roi_end_w = (scalar_t)(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
+    scalar_t roi_end_h = (scalar_t)(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
+
+    // Force too small ROIs to be 1x1
+    scalar_t roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
+    scalar_t roi_height = max(roi_end_h - roi_start_h, 0.1);
+
+    // Compute w and h at bottom
+    scalar_t bin_size_h = roi_height / (scalar_t)(pooled_height);
+    scalar_t bin_size_w = roi_width / (scalar_t)(pooled_width);
+
+    scalar_t sub_bin_size_h = bin_size_h / (scalar_t)(sample_per_part);
+    scalar_t sub_bin_size_w = bin_size_w / (scalar_t)(sample_per_part);
+
+    int part_h = floor((scalar_t)(ph) / pooled_height * part_size);
+    int part_w = floor((scalar_t)(pw) / pooled_width * part_size);
+    int class_id = ctop / channels_each_class;
+    scalar_t trans_x = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+    scalar_t trans_y = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+
+    scalar_t wstart = (scalar_t)(pw)*bin_size_w + roi_start_w;
+    wstart += trans_x * roi_width;
+    scalar_t hstart = (scalar_t)(ph)*bin_size_h + roi_start_h;
+    hstart += trans_y * roi_height;
+
+    scalar_t sum = 0;
+    int count = 0;
+    int gw = floor((scalar_t)(pw)*group_size / pooled_width);
+    int gh = floor((scalar_t)(ph)*group_size / pooled_height);
+    gw = min(max(gw, 0), group_size - 1);
+    gh = min(max(gh, 0), group_size - 1);
+
+    const scalar_t *offset_bottom_data = bottom_data + (roi_batch_ind * channels) * height * width;
+    for (int ih = 0; ih < sample_per_part; ih++)
+    {
+      for (int iw = 0; iw < sample_per_part; iw++)
+      {
+        scalar_t w = wstart + iw * sub_bin_size_w;
+        scalar_t h = hstart + ih * sub_bin_size_h;
+        // bilinear interpolation
+        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
+        {
+          continue;
+        }
+        w = min(max(w, 0.), width - 1.);
+        h = min(max(h, 0.), height - 1.);
+        int c = (ctop * group_size + gh) * group_size + gw;
+        scalar_t val = bilinear_interp(offset_bottom_data + c * height * width, w, h, width, height);
+        sum += val;
+        count++;
+      }
+    }
+    top_data[index] = count == 0 ? (scalar_t)(0) : sum / count;
+    top_count[index] = count;
+  }
+}
+
+template <typename scalar_t>
+__global__ void DeformablePSROIPoolBackwardAccKernel(
+    const int count,
+    const scalar_t *top_diff,
+    const scalar_t *top_count,
+    const int num_rois,
+    const scalar_t spatial_scale,
+    const int channels,
+    const int height, const int width,
+    const int pooled_height, const int pooled_width,
+    const int output_dim,
+    scalar_t *bottom_data_diff, scalar_t *bottom_trans_diff,
+    const scalar_t *bottom_data,
+    const scalar_t *bottom_rois,
+    const scalar_t *bottom_trans,
+    const int no_trans,
+    const scalar_t trans_std,
+    const int sample_per_part,
+    const int group_size,
+    const int part_size,
+    const int num_classes,
+    const int channels_each_class)
+{
+  CUDA_KERNEL_LOOP(index, count)
+  {
+    // The output is in order (n, ctop, ph, pw)
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int ctop = (index / pooled_width / pooled_height) % output_dim;
+    int n = index / pooled_width / pooled_height / output_dim;
+
+    // [start, end) interval for spatial sampling
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 5;
+    int roi_batch_ind = offset_bottom_rois[0];
+    scalar_t roi_start_w = (scalar_t)(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
+    scalar_t roi_start_h = (scalar_t)(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
+    scalar_t roi_end_w = (scalar_t)(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
+    scalar_t roi_end_h = (scalar_t)(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
+
+    // Force too small ROIs to be 1x1
+    scalar_t roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
+    scalar_t roi_height = max(roi_end_h - roi_start_h, 0.1);
+
+    // Compute w and h at bottom
+    scalar_t bin_size_h = roi_height / (scalar_t)(pooled_height);
+    scalar_t bin_size_w = roi_width / (scalar_t)(pooled_width);
+
+    scalar_t sub_bin_size_h = bin_size_h / (scalar_t)(sample_per_part);
+    scalar_t sub_bin_size_w = bin_size_w / (scalar_t)(sample_per_part);
+
+    int part_h = floor((scalar_t)(ph) / pooled_height * part_size);
+    int part_w = floor((scalar_t)(pw) / pooled_width * part_size);
+    int class_id = ctop / channels_each_class;
+    scalar_t trans_x = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+    scalar_t trans_y = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+
+    scalar_t wstart = (scalar_t)(pw)*bin_size_w + roi_start_w;
+    wstart += trans_x * roi_width;
+    scalar_t hstart = (scalar_t)(ph)*bin_size_h + roi_start_h;
+    hstart += trans_y * roi_height;
+
+    if (top_count[index] <= 0)
+    {
+      continue;
+    }
+    scalar_t diff_val = top_diff[index] / top_count[index];
+    const scalar_t *offset_bottom_data = bottom_data + roi_batch_ind * channels * height * width;
+    scalar_t *offset_bottom_data_diff = bottom_data_diff + roi_batch_ind * channels * height * width;
+    int gw = floor((scalar_t)(pw)*group_size / pooled_width);
+    int gh = floor((scalar_t)(ph)*group_size / pooled_height);
+    gw = min(max(gw, 0), group_size - 1);
+    gh = min(max(gh, 0), group_size - 1);
+
+    for (int ih = 0; ih < sample_per_part; ih++)
+    {
+      for (int iw = 0; iw < sample_per_part; iw++)
+      {
+        scalar_t w = wstart + iw * sub_bin_size_w;
+        scalar_t h = hstart + ih * sub_bin_size_h;
+        // bilinear interpolation
+        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
+        {
+          continue;
+        }
+        w = min(max(w, 0.), width - 1.);
+        h = min(max(h, 0.), height - 1.);
+        int c = (ctop * group_size + gh) * group_size + gw;
+        // backward on feature
+        int x0 = floor(w);
+        int x1 = ceil(w);
+        int y0 = floor(h);
+        int y1 = ceil(h);
+        scalar_t dist_x = w - x0, dist_y = h - y0;
+        scalar_t q00 = (1 - dist_x) * (1 - dist_y);
+        scalar_t q01 = (1 - dist_x) * dist_y;
+        scalar_t q10 = dist_x * (1 - dist_y);
+        scalar_t q11 = dist_x * dist_y;
+        int bottom_index_base = c * height * width;
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x0, q00 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x0, q01 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x1, q10 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x1, q11 * diff_val);
+
+        if (no_trans)
+        {
+          continue;
+        }
+        scalar_t U00 = offset_bottom_data[bottom_index_base + y0 * width + x0];
+        scalar_t U01 = offset_bottom_data[bottom_index_base + y1 * width + x0];
+        scalar_t U10 = offset_bottom_data[bottom_index_base + y0 * width + x1];
+        scalar_t U11 = offset_bottom_data[bottom_index_base + y1 * width + x1];
+        scalar_t diff_x = (U11 * dist_y + U10 * (1 - dist_y) - U01 * dist_y - U00 * (1 - dist_y)) * trans_std * diff_val;
+        diff_x *= roi_width;
+        scalar_t diff_y = (U11 * dist_x + U01 * (1 - dist_x) - U10 * dist_x - U00 * (1 - dist_x)) * trans_std * diff_val;
+        diff_y *= roi_height;
+
+        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w, diff_x);
+        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w, diff_y);
+      }
+    }
+  }
+}
+
+void DeformablePSROIPoolForward(const at::Tensor data,
+                                const at::Tensor bbox,
+                                const at::Tensor trans,
+                                at::Tensor out,
+                                at::Tensor top_count,
+                                const int batch,
+                                const int channels,
+                                const int height,
+                                const int width,
+                                const int num_bbox,
+                                const int channels_trans,
+                                const int no_trans,
+                                const float spatial_scale,
+                                const int output_dim,
+                                const int group_size,
+                                const int pooled_size,
+                                const int part_size,
+                                const int sample_per_part,
+                                const float trans_std)
+{
+  const int pooled_height = pooled_size;
+  const int pooled_width = pooled_size;
+  const int count = num_bbox * output_dim * pooled_height * pooled_width;
+  const int num_classes = no_trans ? 1 : channels_trans / 2;
+  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data.type(), "deformable_psroi_pool_forward", ([&] {
+        const scalar_t *bottom_data = data.data<scalar_t>();
+        const scalar_t *bottom_rois = bbox.data<scalar_t>();
+        const scalar_t *bottom_trans = no_trans ? NULL : trans.data<scalar_t>();
+        scalar_t *top_data = out.data<scalar_t>();
+        scalar_t *top_count_data = top_count.data<scalar_t>();
+
+        DeformablePSROIPoolForwardKernel<<<GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
+            count, bottom_data, (scalar_t)spatial_scale, channels, height, width, pooled_height, pooled_width,
+            bottom_rois, bottom_trans, no_trans, (scalar_t)trans_std, sample_per_part, output_dim,
+            group_size, part_size, num_classes, channels_each_class, top_data, top_count_data);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in DeformablePSROIPoolForward: %s\n", cudaGetErrorString(err));
+  }
+}
+
+void DeformablePSROIPoolBackwardAcc(const at::Tensor out_grad,
+                                    const at::Tensor data,
+                                    const at::Tensor bbox,
+                                    const at::Tensor trans,
+                                    const at::Tensor top_count,
+                                    at::Tensor in_grad,
+                                    at::Tensor trans_grad,
+                                    const int batch,
+                                    const int channels,
+                                    const int height,
+                                    const int width,
+                                    const int num_bbox,
+                                    const int channels_trans,
+                                    const int no_trans,
+                                    const float spatial_scale,
+                                    const int output_dim,
+                                    const int group_size,
+                                    const int pooled_size,
+                                    const int part_size,
+                                    const int sample_per_part,
+                                    const float trans_std)
+{
+  // LOG(INFO) << "DeformablePSROIPoolBackward";
+  const int num_rois = num_bbox;
+  const int pooled_height = pooled_size;
+  const int pooled_width = pooled_size;
+  const int count = num_bbox * output_dim * pooled_height * pooled_width;
+  const int num_classes = no_trans ? 1 : channels_trans / 2;
+  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      out_grad.type(), "deformable_psroi_pool_backward_acc", ([&] {
+        const scalar_t *top_diff = out_grad.data<scalar_t>();
+        const scalar_t *bottom_data = data.data<scalar_t>();
+        const scalar_t *bottom_rois = bbox.data<scalar_t>();
+        const scalar_t *bottom_trans = no_trans ? NULL : trans.data<scalar_t>();
+        scalar_t *bottom_data_diff = in_grad.data<scalar_t>();
+        scalar_t *bottom_trans_diff = no_trans ? NULL : trans_grad.data<scalar_t>();
+        const scalar_t *top_count_data = top_count.data<scalar_t>();
+
+        DeformablePSROIPoolBackwardAccKernel<<<GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
+            count, top_diff, top_count_data, num_rois, (scalar_t)spatial_scale, channels, height, width,
+            pooled_height, pooled_width, output_dim, bottom_data_diff, bottom_trans_diff,
+            bottom_data, bottom_rois, bottom_trans, no_trans, (scalar_t)trans_std, sample_per_part,
+            group_size, part_size, num_classes, channels_each_class);
+      }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in DeformablePSROIPoolForward: %s\n", cudaGetErrorString(err));
+  }
+}

IndicPhotoOCR/detection/textbpn/network/backbone/resnet.py

@@ -0,0 +1,336 @@
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+BatchNorm2d = nn.BatchNorm2d
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def constant_init(module, constant, bias=0):
+    nn.init.constant_(module.weight, constant)
+    if hasattr(module, 'bias'):
+        nn.init.constant_(module.bias, bias)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None):
+        super(BasicBlock, self).__init__()
+        self.with_dcn = dcn is not None
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.with_modulated_dcn = False
+        if self.with_dcn:
+            fallback_on_stride = dcn.get('fallback_on_stride', False)
+            self.with_modulated_dcn = dcn.get('modulated', False)
+        # self.conv2 = conv3x3(planes, planes)
+        if not self.with_dcn or fallback_on_stride:
+            self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                                   padding=1, bias=False)
+        else:
+            deformable_groups = dcn.get('deformable_groups', 1)
+            if not self.with_modulated_dcn:
+                from network.backbone.assets.dcn import DeformConv
+                conv_op = DeformConv
+                offset_channels = 18
+            else:
+                from network.backbone.assets.dcn import ModulatedDeformConv
+                conv_op = ModulatedDeformConv
+                offset_channels = 27
+            self.conv2_offset = nn.Conv2d(
+                planes,
+                deformable_groups * offset_channels,
+                kernel_size=3,
+                padding=1)
+            self.conv2 = conv_op(
+                planes,
+                planes,
+                kernel_size=3,
+                padding=1,
+                deformable_groups=deformable_groups,
+                bias=False)
+        self.bn2 = BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        # out = self.conv2(out)
+        if not self.with_dcn:
+            out = self.conv2(out)
+        elif self.with_modulated_dcn:
+            offset_mask = self.conv2_offset(out)
+            offset = offset_mask[:, :18, :, :]
+            mask = offset_mask[:, -9:, :, :].sigmoid()
+            out = self.conv2(out, offset, mask)
+        else:
+            offset = self.conv2_offset(out)
+            out = self.conv2(out, offset)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None):
+        super(Bottleneck, self).__init__()
+        self.with_dcn = dcn is not None
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm2d(planes)
+        fallback_on_stride = False
+        self.with_modulated_dcn = False
+        if self.with_dcn:
+            fallback_on_stride = dcn.get('fallback_on_stride', False)
+            self.with_modulated_dcn = dcn.get('modulated', False)
+        if not self.with_dcn or fallback_on_stride:
+            self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                                   stride=stride, padding=1, bias=False)
+        else:
+            deformable_groups = dcn.get('deformable_groups', 1)
+            if not self.with_modulated_dcn:
+                from network.backbone.assets.dcn import DeformConv
+                conv_op = DeformConv
+                offset_channels = 18
+            else:
+                from network.backbone.assets.dcn import ModulatedDeformConv
+                conv_op = ModulatedDeformConv
+                offset_channels = 27
+            self.conv2_offset = nn.Conv2d(
+                planes, deformable_groups * offset_channels,
+                kernel_size=3,
+                padding=1)
+            self.conv2 = conv_op(
+                planes, planes, kernel_size=3, padding=1, stride=stride,
+                deformable_groups=deformable_groups, bias=False)
+        self.bn2 = BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.dcn = dcn
+        self.with_dcn = dcn is not None
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        # out = self.conv2(out)
+        if not self.with_dcn:
+            out = self.conv2(out)
+        elif self.with_modulated_dcn:
+            offset_mask = self.conv2_offset(out)
+            offset = offset_mask[:, :18, :, :]
+            mask = offset_mask[:, -9:, :, :].sigmoid()
+            out = self.conv2(out, offset, mask)
+        else:
+            offset = self.conv2_offset(out)
+            out = self.conv2(out, offset)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000, 
+                 dcn=None, stage_with_dcn=(False, False, False, False)):
+        self.dcn = dcn
+        self.stage_with_dcn = stage_with_dcn
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(
+            block, 128, layers[1], stride=2, dcn=dcn)
+        self.layer3 = self._make_layer(
+            block, 256, layers[2], stride=2, dcn=dcn)
+        self.layer4 = self._make_layer(
+            block, 512, layers[3], stride=2, dcn=dcn)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+    
+        self.smooth = nn.Conv2d(2048, 256, kernel_size=1, stride=1, padding=1)    
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        if self.dcn is not None:
+            for m in self.modules():
+                if isinstance(m, Bottleneck) or isinstance(m, BasicBlock):
+                    if hasattr(m, 'conv2_offset'):
+                        constant_init(m.conv2_offset, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dcn=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes,
+                            stride, downsample, dcn=dcn))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dcn=dcn))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x1 = self.maxpool(x)
+
+        x2 = self.layer1(x1)
+        x3 = self.layer2(x2)
+        x4 = self.layer3(x3)
+        x5 = self.layer4(x4)
+
+        return x1, x2, x3, x4, x5
+
+
+def resnet18(pretrained=True, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet18']), strict=False)
+    return model
+
+def deformable_resnet18(pretrained=True, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2],
+                    dcn=dict(modulated=True,
+                            deformable_groups=1,
+                            fallback_on_stride=False),
+                    stage_with_dcn=[False, True, True, True], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet18']), strict=False)
+    return model
+
+
+def resnet34(pretrained=True, **kwargs):
+    """Constructs a ResNet-34 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet34']), strict=False)
+    return model
+
+
+def resnet50(pretrained=True, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet50']), strict=False)
+    return model
+
+
+def deformable_resnet50(pretrained=True, **kwargs):
+    """Constructs a ResNet-50 model with deformable conv.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3],
+                   dcn=dict(modulated=True,
+                            deformable_groups=1,
+                            fallback_on_stride=False),
+                   stage_with_dcn=[False, True, True, True],
+                   **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet50']), strict=False)
+    return model
+
+
+def resnet101(pretrained=True, **kwargs):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet101']), strict=False)
+    return model
+
+
+def resnet152(pretrained=True, **kwargs):
+    """Constructs a ResNet-152 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(
+            model_urls['resnet152']), strict=False)
+    return model

IndicPhotoOCR/detection/textbpn/network/backbone/vgg.py

@@ -0,0 +1,60 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torchvision.models as models
+
+model_urls = {
+    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
+    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
+    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
+    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
+    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
+    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
+}
+
+
+class VggNet(nn.Module):
+    def __init__(self, name="vgg16", pretrain=True):
+        super().__init__()
+        if name == "vgg16":
+            base_net = models.vgg16(pretrained=False)
+        elif name == "vgg16_bn":
+            base_net = models.vgg16_bn(pretrained=False)
+        else:
+            print(" base model is not support !")
+        if pretrain:
+            print("load the {} weight from ./cache".format(name))
+            base_net.load_state_dict(model_zoo.load_url(model_urls[name], model_dir="./cache"))
+
+        if name == "vgg16":
+            self.stage1 = nn.Sequential(*[base_net.features[layer] for layer in range(0, 5)])
+            self.stage2 = nn.Sequential(*[base_net.features[layer] for layer in range(5, 10)])
+            self.stage3 = nn.Sequential(*[base_net.features[layer] for layer in range(10, 17)])
+            self.stage4 = nn.Sequential(*[base_net.features[layer] for layer in range(17, 24)])
+            self.stage5 = nn.Sequential(*[base_net.features[layer] for layer in range(24, 31)])
+        elif name == "vgg16_bn":
+            self.stage1 = nn.Sequential(*[base_net.features[layer] for layer in range(0, 7)])
+            self.stage2 = nn.Sequential(*[base_net.features[layer] for layer in range(7, 14)])
+            self.stage3 = nn.Sequential(*[base_net.features[layer] for layer in range(14, 24)])
+            self.stage4 = nn.Sequential(*[base_net.features[layer] for layer in range(24, 34)])
+            self.stage5 = nn.Sequential(*[base_net.features[layer] for layer in range(34, 44)])
+
+    def forward(self, x):
+        C1 = self.stage1(x)
+        C2 = self.stage2(C1)
+        C3 = self.stage3(C2)
+        C4 = self.stage4(C3)
+        C5 = self.stage5(C4)
+
+        return C1, C2, C3, C4, C5
+
+
+if __name__ == '__main__':
+    import torch
+    input = torch.randn((4, 3, 512, 512))
+    net = VggNet()
+    C1, C2, C3, C4, C5 = net(input)
+    print(C1.size())
+    print(C2.size())
+    print(C3.size())
+    print(C4.size())
+    print(C5.size())

IndicPhotoOCR/detection/textbpn/network/layers/Adaptive_Deformation.py

@@ -0,0 +1,88 @@
+###################################################################
+# File Name: AdaptiveDeformation.py
+# Author: S.X.Zhang
+###################################################################
+
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+
+class MeanAggregator(nn.Module):
+    def __init__(self):
+        super(MeanAggregator, self).__init__()
+
+    def forward(self, features, A):
+        x = torch.bmm(A, features)
+        return x
+
+
+class GraphConv(nn.Module):
+    def __init__(self, in_dim, out_dim, agg):
+        super(GraphConv, self).__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.weight = nn.Parameter(torch.FloatTensor(in_dim * 2, out_dim))
+        self.bias = nn.Parameter(torch.FloatTensor(out_dim))
+        init.xavier_uniform_(self.weight)
+        init.constant_(self.bias, 0)
+        self.agg = agg()
+
+    def forward(self, features, A):
+        b, n, d = features.shape
+        assert (d == self.in_dim)
+        agg_feats = self.agg(features, A)
+        cat_feats = torch.cat([features, agg_feats], dim=2)
+        out = torch.einsum('bnd,df->bnf', (cat_feats, self.weight))
+        out = F.relu(out + self.bias)
+        return out
+
+
+class AdaptiveDeformation(nn.Module):
+    def __init__(self, input, state_dim):
+        super(AdaptiveDeformation, self).__init__()
+        self.bn0 = nn.BatchNorm1d(input, affine=False)
+        self.conv1 = nn.Conv1d(input, state_dim, 1)
+        self.rnn = nn.LSTM(input, state_dim, 1, bidirectional=True)
+        self.gconv1 = GraphConv(input, 256, MeanAggregator)
+        self.gconv2 = GraphConv(256, 1024, MeanAggregator)
+        self.gconv3 = GraphConv(1024, 512, MeanAggregator)
+        self.gconv4 = GraphConv(512, state_dim, MeanAggregator)
+
+        self.prediction = nn.Sequential(
+            nn.Conv1d(4*state_dim, 128, 1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv1d(128, 64, 1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv1d(64, 2, 1))
+
+    def forward(self, x, A):
+        x = self.bn0(x)
+
+        # # rnn block
+        yl = x.permute(2, 0, 1)
+        yl, _ = self.rnn(yl)
+        yl = yl.permute(1, 2, 0)
+
+        # # gcn block
+        yg = x.permute(0, 2, 1)
+        b, n, c = yg.shape
+        A = A.expand(b, n, n)
+        yg = self.gconv1(yg, A)
+        yg = self.gconv2(yg, A)
+        yg = self.gconv3(yg, A)
+        yg = self.gconv4(yg, A)
+        yg = yg.permute(0, 2, 1)
+
+        # res block
+        x = torch.cat([yl, yg, self.conv1(x)], dim=1)
+        pred = self.prediction(x)
+
+        return pred

IndicPhotoOCR/detection/textbpn/network/layers/CircConv.py

@@ -0,0 +1,91 @@
+import torch.nn as nn
+import torch
+
+
+class CircConv(nn.Module):
+    def __init__(self, state_dim, out_state_dim=None, n_adj=4):
+        super(CircConv, self).__init__()
+
+        self.n_adj = n_adj
+        out_state_dim = state_dim if out_state_dim is None else out_state_dim
+        self.fc = nn.Conv1d(state_dim, out_state_dim, kernel_size=self.n_adj*2+1)
+
+    def forward(self, input, adj):
+        input = torch.cat([input[..., -self.n_adj:], input, input[..., :self.n_adj]], dim=2)
+        return self.fc(input)
+
+
+class DilatedCircConv(nn.Module):
+    def __init__(self, state_dim, out_state_dim=None, n_adj=4, dilation=1):
+        super(DilatedCircConv, self).__init__()
+
+        self.n_adj = n_adj
+        self.dilation = dilation
+        out_state_dim = state_dim if out_state_dim is None else out_state_dim
+        self.fc = nn.Conv1d(state_dim, out_state_dim, kernel_size=self.n_adj*2+1, dilation=self.dilation)
+
+    def forward(self, input, adj):
+        if self.n_adj != 0:
+            input = torch.cat([input[..., -self.n_adj*self.dilation:], input, input[..., :self.n_adj*self.dilation]], dim=2)
+        return self.fc(input)
+
+
+_conv_factory = {
+    'grid': CircConv,
+    'dgrid': DilatedCircConv
+}
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, state_dim, out_state_dim, conv_type, n_adj=4, dilation=1):
+        super(BasicBlock, self).__init__()
+
+        self.conv = _conv_factory[conv_type](state_dim, out_state_dim, n_adj, dilation)
+        self.relu = nn.ReLU(inplace=True)
+        self.norm = nn.BatchNorm1d(out_state_dim)
+
+    def forward(self, x, adj=None):
+        x = self.conv(x, adj)
+        x = self.relu(x)
+        x = self.norm(x)
+        return x
+
+
+class DeepSnake(nn.Module):
+    def __init__(self, state_dim, feature_dim, conv_type='dgrid'):
+        super(DeepSnake, self).__init__()
+
+        self.head = BasicBlock(feature_dim, state_dim, conv_type)
+
+        self.res_layer_num = 7
+        dilation = [1, 1, 1, 2, 2, 4, 4]
+        for i in range(self.res_layer_num):
+            conv = BasicBlock(state_dim, state_dim, conv_type, n_adj=4, dilation=dilation[i])
+            self.__setattr__('res'+str(i), conv)
+
+        fusion_state_dim = 256
+        self.fusion = nn.Conv1d(state_dim * (self.res_layer_num + 1), fusion_state_dim, 1)
+        self.prediction = nn.Sequential(
+            nn.Conv1d(state_dim * (self.res_layer_num + 1) + fusion_state_dim, 256, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(256, 64, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(64, 2, 1)
+        )
+
+    def forward(self, x, adj):
+        states = []
+
+        x = self.head(x, adj)
+        states.append(x)
+        for i in range(self.res_layer_num):
+            x = self.__getattr__('res'+str(i))(x, adj) + x
+            states.append(x)
+
+        state = torch.cat(states, dim=1)
+        global_state = torch.max(self.fusion(state), dim=2, keepdim=True)[0]
+        global_state = global_state.expand(global_state.size(0), global_state.size(1), state.size(2))
+        state = torch.cat([global_state, state], dim=1)
+        x = self.prediction(state)
+
+        return x

IndicPhotoOCR/detection/textbpn/network/layers/GCN.py

@@ -0,0 +1,77 @@
+###################################################################
+# File Name: GCN.py
+# Author: S.X.Zhang
+###################################################################
+
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+
+class MeanAggregator(nn.Module):
+    def __init__(self):
+        super(MeanAggregator, self).__init__()
+
+    def forward(self, features, A):
+        x = torch.bmm(A, features)
+        return x
+
+
+class GraphConv(nn.Module):
+    def __init__(self, in_dim, out_dim, agg):
+        super(GraphConv, self).__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.weight = nn.Parameter(torch.FloatTensor(in_dim * 2, out_dim))
+        self.bias = nn.Parameter(torch.FloatTensor(out_dim))
+        init.xavier_uniform_(self.weight)
+        init.constant_(self.bias, 0)
+        self.agg = agg()
+
+    def forward(self, features, A):
+        b, n, d = features.shape
+        assert (d == self.in_dim)
+        agg_feats = self.agg(features, A)
+        cat_feats = torch.cat([features, agg_feats], dim=2)
+        out = torch.einsum('bnd,df->bnf', (cat_feats, self.weight))
+        out = F.relu(out + self.bias)
+        return out
+
+
+class GCN(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super(GCN, self).__init__()
+        self.bn0 = nn.BatchNorm1d(in_dim, affine=False)
+
+        self.conv1 = GraphConv(in_dim, 256, MeanAggregator)
+        self.conv2 = GraphConv(256, 1024, MeanAggregator)
+        self.conv3 = GraphConv(1024, 512, MeanAggregator)
+        self.conv4 = GraphConv(512, out_dim, MeanAggregator)
+
+        self.prediction = nn.Sequential(
+            nn.Conv1d(out_dim, 128, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(128, 64, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(64, 2, 1))
+
+    def forward(self, x, A):
+        x = self.bn0(x)
+        x = x.permute(0, 2, 1)
+        b, n, c = x.shape
+        A = A.expand(b, n, n)
+
+        x = self.conv1(x, A)
+        x = self.conv2(x, A)
+        x = self.conv3(x, A)
+        x = self.conv4(x, A)
+
+        x = x.permute(0, 2, 1)
+        pred = self.prediction(x)
+
+        return pred

IndicPhotoOCR/detection/textbpn/network/layers/GraphConv.py

@@ -0,0 +1,45 @@
+import math
+
+import torch
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from torch.nn import init
+
+
+class GraphConvolution(Module):
+    """
+    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
+    """
+
+    def __init__(self, in_features, out_features, bias=True):
+        super(GraphConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        init.xavier_uniform_(self.weight)
+        if bias:
+            self.bias = Parameter(torch.FloatTensor(out_features))
+            init.constant_(self.bias, 0)
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, adj):
+        support = torch.mm(input, self.weight)
+        output = torch.spmm(adj, support)
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'

IndicPhotoOCR/detection/textbpn/network/layers/RNN.py

@@ -0,0 +1,35 @@
+###################################################################
+# File Name: RNN.py
+# Author: S.X.Zhang
+###################################################################
+
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+
+class RNN(nn.Module):
+    def __init__(self, input, state_dim):
+        super(RNN, self).__init__()
+        self.bn0 = nn.BatchNorm1d(input, affine=False)
+        self.rnn = nn.LSTM(input, state_dim, 1, dropout=0.1, bidirectional=True)
+        self.prediction = nn.Sequential(
+            nn.Conv1d(state_dim*2, 128, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(128, 64, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(64, 2, 1))
+
+    def forward(self, x, adj):
+        x = self.bn0(x)
+        x = x.permute(2, 0, 1)
+        x, _ = self.rnn(x)
+        x = x.permute(1, 2, 0)
+        pred = self.prediction(x)
+
+        return pred

IndicPhotoOCR/detection/textbpn/network/layers/Transformer.py

@@ -0,0 +1,140 @@
+###################################################################
+# File Name: GCN.py
+# Author: S.X.Zhang
+###################################################################
+import torch
+from torch import nn, Tensor
+import numpy as np
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+
+
+class Positional_encoding(nn.Module):
+    def __init__(self, PE_size, n_position=256):
+        super(Positional_encoding, self).__init__()
+        self.PE_size = PE_size
+        self.n_position = n_position
+        self.register_buffer('pos_table', self.get_encoding_table(n_position, PE_size))
+
+    def get_encoding_table(self, n_position, PE_size):
+        position_table = np.array(
+            [[pos / np.power(10000, 2. * i / self.PE_size) for i in range(self.PE_size)] for pos in range(n_position)])
+        position_table[:, 0::2] = np.sin(position_table[:, 0::2])
+        position_table[:, 1::2] = np.cos(position_table[:, 1::2])
+        return torch.FloatTensor(position_table).unsqueeze(0)
+
+    def forward(self, inputs):
+        return inputs + self.pos_table[:, :inputs.size(1), :].clone().detach()
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, num_heads, embed_dim, dropout=0.1, if_resi=True):
+        super(MultiHeadAttention, self).__init__()
+        self.layer_norm = nn.LayerNorm(embed_dim)
+        self.MultiheadAttention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
+        self.Q_proj = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.ReLU())
+        self.K_proj = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.ReLU())
+        self.V_proj = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.ReLU())
+        self.if_resi = if_resi
+
+    def forward(self, inputs):
+        query = self.layer_norm(inputs)
+        q = self.Q_proj(query)
+        k = self.K_proj(query)
+        v = self.V_proj(query)
+        attn_output, attn_output_weights = self.MultiheadAttention(q, k, v)
+        if self.if_resi:
+            attn_output += inputs
+        else:
+            attn_output = attn_output
+
+        return attn_output
+
+
+class FeedForward(nn.Module):
+    def __init__(self, in_channel, FFN_channel, if_resi=True):
+        super(FeedForward, self).__init__()
+        """
+        1024 2048
+        """
+        output_channel = (FFN_channel, in_channel)
+        self.fc1 = nn.Sequential(nn.Linear(in_channel, output_channel[0]), nn.ReLU())
+        self.fc2 = nn.Linear(output_channel[0], output_channel[1])
+        self.layer_norm = nn.LayerNorm(in_channel)
+        self.if_resi = if_resi
+
+    def forward(self, inputs):
+        outputs = self.layer_norm(inputs)
+        outputs = self.fc1(outputs)
+        outputs = self.fc2(outputs)
+        if self.if_resi:
+            outputs += inputs
+        else:
+            outputs = outputs
+        return outputs
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, out_dim, in_dim, num_heads, attention_size,
+                 dim_feedforward=1024, drop_rate=0.1, if_resi=True, block_nums=3):
+        super(TransformerLayer, self).__init__()
+        self.block_nums = block_nums
+        self.if_resi = if_resi
+        self.linear = nn.Linear(in_dim, attention_size)
+        for i in range(self.block_nums):
+            self.__setattr__('MHA_self_%d' % i, MultiHeadAttention(num_heads, attention_size,
+                                                                   dropout=drop_rate, if_resi=if_resi))
+            self.__setattr__('FFN_%d' % i, FeedForward(out_dim, dim_feedforward, if_resi=if_resi))
+
+    def forward(self, query):
+        inputs = self.linear(query)
+        # outputs = inputs
+        for i in range(self.block_nums):
+            outputs = self.__getattr__('MHA_self_%d' % i)(inputs)
+            outputs = self.__getattr__('FFN_%d' % i)(outputs)
+            if self.if_resi:
+                inputs = inputs+outputs
+            else:
+                inputs = outputs
+        # outputs = inputs
+        return inputs
+
+
+class Transformer(nn.Module):
+
+    def __init__(self, in_dim, out_dim, num_heads=8,
+                 dim_feedforward=1024, drop_rate=0.1, if_resi=False, block_nums=3):
+        super().__init__()
+
+        self.bn0 = nn.BatchNorm1d(in_dim, affine=False)
+        self.conv1 = nn.Conv1d(in_dim, out_dim, 1, dilation=1)
+
+        # self.pos_embedding = Positional_encoding(in_dim)
+        self.transformer = TransformerLayer(out_dim, in_dim, num_heads, attention_size=out_dim,
+                                            dim_feedforward=dim_feedforward, drop_rate=drop_rate,
+                                            if_resi=if_resi, block_nums=block_nums)
+
+        self.prediction = nn.Sequential(
+            nn.Conv1d(2*out_dim, 128, 1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv1d(128, 64, 1),
+            nn.ReLU(inplace=True),
+            # nn.Dropout(0.1),
+            nn.Conv1d(64, 2, 1))
+
+    def forward(self, x, adj):
+        x = self.bn0(x)
+
+        x1 = x.permute(0, 2, 1)
+        # x1 = self.pos_embedding(x1)
+        x1 = self.transformer(x1)
+        x1 = x1.permute(0, 2, 1)
+
+        x = torch.cat([x1, self.conv1(x)], dim=1)
+        # x = x1+self.conv1(x)
+        pred = self.prediction(x)
+
+        return pred
+
+
+

IndicPhotoOCR/detection/textbpn/network/layers/Transformer_old.py

@@ -0,0 +1,171 @@
+###################################################################
+# File Name: GCN.py
+# Author: S.X.Zhang
+###################################################################
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch.autograd import Variable
+import numpy as np
+from cfglib.config import config as cfg
+
+
+class Positional_encoding(nn.Module):
+    def __init__(self, PE_size, n_position=200):
+        super(Positional_encoding, self).__init__()
+        self.PE_size = PE_size
+        self.n_position = n_position
+        self.register_buffer('pos_table', self.get_encoding_table(n_position, PE_size))
+
+    def get_encoding_table(self, n_position, PE_size):
+        position_table = np.array(
+            [[pos / np.power(10000, 2. * i / self.PE_size) for i in range(self.PE_size)] for pos in range(n_position)])
+        position_table[:, 0::2] = np.sin(position_table[:, 0::2])
+        position_table[:, 1::2] = np.cos(position_table[:, 1::2])
+        return torch.FloatTensor(position_table).unsqueeze(0)
+
+    def forward(self, inputs):
+        return inputs + self.pos_table[:, :inputs.size(1), :].clone().detach()
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, num_heads, embedding_size, attention_size,
+                 drop_rate, future_blind=True, query_mask=False, if_resi=True):
+        super(MultiHeadAttention, self).__init__()
+        self.num_heads = num_heads
+        self.embedding_size = embedding_size
+        self.attention_size = attention_size
+        self.drop_rate = drop_rate
+        self.future_blind = future_blind
+        
+        self.Q_proj = nn.Sequential(nn.Linear(self.embedding_size, self.attention_size), nn.ReLU())
+        self.K_proj = nn.Sequential(nn.Linear(self.embedding_size, self.attention_size), nn.ReLU())
+        self.V_proj = nn.Sequential(nn.Linear(self.embedding_size, self.attention_size), nn.ReLU())
+
+        self.drop_out = nn.Dropout(p=self.drop_rate)
+        self.layer_norm = nn.LayerNorm(self.attention_size)
+        self.if_resi = if_resi
+
+    def forward(self, query, key, value):
+        q = self.Q_proj(query)
+        k = self.K_proj(key)
+        v = self.V_proj(value)
+
+        q_ = torch.cat(torch.chunk(q, self.num_heads, dim=2), dim=0)
+        k_ = torch.cat(torch.chunk(k, self.num_heads, dim=2), dim=0)
+        v_ = torch.cat(torch.chunk(v, self.num_heads, dim=2), dim=0)
+
+        outputs = torch.bmm(q_, k_.permute(0, 2, 1))
+        outputs = outputs / (k_.size()[-1] ** 0.5)
+
+        # key mask
+
+        # future mask
+        if self.future_blind:
+            diag_vals = torch.ones_like(outputs[0, :, :]).to(cfg.device)
+            tril = torch.tril(diag_vals, diagonal=0)
+            masks = Variable(torch.unsqueeze(tril, 0).repeat(outputs.size()[0], 1, 1))  # (h*N,T_q,T_k)
+            padding = Variable(torch.ones_like(masks).to(cfg.device) * (-2 ** 32 + 1))
+            condition = masks.eq(0)
+            outputs = torch.where(condition, padding, outputs)
+
+        outputs = F.softmax(outputs, dim=-1)
+        # if self.future_blind==True:a
+        #     print(outputs[0])
+        outputs = self.drop_out(outputs)
+
+        outputs = torch.bmm(outputs, v_)
+        outputs = torch.cat(torch.chunk(outputs, self.num_heads, dim=0), dim=2)  # N,T_q,C
+
+        if self.if_resi:
+            # outputs += query
+            outputs += q
+        else:
+            outputs = outputs
+        outputs = self.layer_norm(outputs)
+
+        return outputs
+
+
+class FeedForward(nn.Module):
+    def __init__(self, in_channel, FFN_channel, if_resi=True):
+        super(FeedForward, self).__init__()
+        """
+        1024 2048
+        """
+        output_channel = (FFN_channel, in_channel)
+        self.fc1 = nn.Sequential(nn.Linear(in_channel, output_channel[0]), nn.ReLU())
+        self.fc2 = nn.Linear(output_channel[0], output_channel[1])
+        self.layer_norm = nn.LayerNorm(in_channel)
+        self.if_resi = if_resi
+
+    def forward(self, inputs):
+        outputs = self.fc1(inputs)
+        outputs = self.fc2(outputs)
+        if self.if_resi:
+            outputs += inputs
+        else:
+            outputs = outputs
+        outputs = self.layer_norm(outputs)
+        return outputs
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, out_dim, num_heads, embedding_size, attention_size,
+                 dim_feedforward=1024, drop_rate=0.1, if_resi=True, block_nums=3):
+        super(TransformerLayer, self).__init__()
+        self.block_nums = block_nums
+        self.if_resi = if_resi
+        for i in range(self.block_nums):
+            self.__setattr__('MHA_self_%d' % i, MultiHeadAttention(num_heads, embedding_size, attention_size,
+                                                                   drop_rate, future_blind=False, if_resi=if_resi))
+            self.__setattr__('FFN_%d' % i, FeedForward(out_dim, dim_feedforward, if_resi=if_resi))
+
+    def forward(self, query):
+        outputs = None
+        for i in range(self.block_nums):
+            outputs = self.__getattr__('MHA_self_%d' % i)(query, query, query)
+            outputs = self.__getattr__('FFN_%d' % i)(outputs)
+        return outputs
+
+
+class Transformer(nn.Module):
+
+    def __init__(self, in_dim, out_dim, num_heads=8,
+                 dim_feedforward=1024, drop_rate=0.1, if_resi=False, block_nums=3):
+        super().__init__()
+
+        self.bn0 = nn.BatchNorm1d(in_dim, affine=False)
+        self.conv1 = nn.Conv1d(in_dim, out_dim, 1, dilation=1)
+
+        embed_dim = in_dim
+        # self.pos_embedding = Positional_encoding(embed_dim)
+        self.transformer = TransformerLayer(out_dim, num_heads, embedding_size=embed_dim,
+                                            attention_size=out_dim, dim_feedforward=dim_feedforward,
+                                            drop_rate=drop_rate, if_resi=if_resi, block_nums=block_nums)
+
+        self.prediction = nn.Sequential(
+            nn.Conv1d(out_dim*2, 128, 1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv1d(128, 64, 1),
+            nn.ReLU(inplace=True),
+            # nn.Dropout(0.1),
+            nn.Conv1d(64, 2, 1))
+
+    def forward(self, x, adj):
+        x = self.bn0(x)
+
+        x1 = x.permute(0, 2, 1)
+        x1 = self.transformer(x1)
+        x1 = x1.permute(0, 2, 1)
+
+        x = torch.cat([x1, self.conv1(x)], dim=1)
+        # x = x1+self.conv1(x)
+        pred = self.prediction(x)
+
+        return pred
+
+
+

IndicPhotoOCR/detection/textbpn/network/layers/gcn_utils.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+__author__ = "S.X.Zhang"
+import torch
+import numpy as np
+import cv2
+import torch.nn as nn
+from torch.autograd import Variable
+
+
+def normalize_adj(A, type="AD"):
+    if type == "DAD":
+        A = A + np.eye(A.shape[0])  # A=A+I
+        d = np.sum(A, axis=0)
+        d_inv = np.power(d, -0.5).flatten()
+        d_inv[np.isinf(d_inv)] = 0.0
+        d_inv = np.diag(d_inv)
+        G = A.dot(d_inv).transpose().dot(d_inv)  # L = D^-1/2 A D^-1/2
+        G = torch.from_numpy(G)
+    elif type == "AD":
+        A = A + np.eye(A.shape[0])  # A=A+I
+        A = torch.from_numpy(A)
+        D = A.sum(1, keepdim=True)
+        G = A.div(D)  # L= A/D
+    else:
+        A = A + np.eye(A.shape[0])  # A=A+I
+        D = A.sum(1, keepdim=True)
+        D = np.diag(D)
+        G = torch.from_numpy(D - A)  # L = D-A
+    return G
+
+
+def np_to_variable(x, is_cuda=True, dtype=torch.FloatTensor):
+    v = Variable(torch.from_numpy(x).type(dtype))
+    if is_cuda:
+        v = v.cuda()
+    return v
+
+
+def set_trainable(model, requires_grad):
+    for param in model.parameters():
+        param.requires_grad = requires_grad
+
+
+def weights_normal_init(model, dev=0.01):
+    if isinstance(model, list):
+        for m in model:
+            weights_normal_init(m, dev)
+    else:
+        for m in model.modules():
+            if isinstance(m, nn.Conv2d):
+                m.weight.data.normal_(0.0, dev)
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0.0, dev)
+
+
+def clip_gradient(model, clip_norm):
+    """Computes a gradient clipping coefficient based on gradient norm."""
+    totalnorm = 0
+    for p in model.parameters():
+        if p.requires_grad:
+            modulenorm = p.grad.data.norm()
+            totalnorm += modulenorm ** 2
+    totalnorm = np.sqrt(totalnorm)
+
+    norm = clip_norm / max(totalnorm, clip_norm)
+    for p in model.parameters():
+        if p.requires_grad:
+            p.grad.mul_(norm)
+
+
+def EuclideanDistances(A, B):
+    BT = B.transpose()
+    vecProd = np.dot(A,BT)
+    SqA = A**2
+    sumSqA = np.matrix(np.sum(SqA, axis=1))
+    sumSqAEx = np.tile(sumSqA.transpose(), (1, vecProd.shape[1]))
+
+    SqB = B**2
+    sumSqB = np.sum(SqB, axis=1)
+    sumSqBEx = np.tile(sumSqB, (vecProd.shape[0], 1))
+    SqED = sumSqBEx + sumSqAEx - 2*vecProd
+    SqED[SqED<0]=0.0
+    ED = np.sqrt(SqED)
+    return ED
+
+
+def get_center_feature(cnn_feature, img_poly, ind, h, w):
+    batch_size = cnn_feature.size(0)
+    for i in range(batch_size):
+        poly = img_poly[ind == i].cpu().numpy()
+        mask = np.zeros((h, w), dtype=np.uint8)
+        cv2.fillPoly(mask, poly.astype(np.int32), color=(1,))
+    return None
+
+
+def get_node_feature(cnn_feature, img_poly, ind, h, w):
+    img_poly = img_poly.clone().float()
+    img_poly[..., 0] = img_poly[..., 0] / (w / 2.) - 1
+    img_poly[..., 1] = img_poly[..., 1] / (h / 2.) - 1
+
+    batch_size = cnn_feature.size(0)
+    gcn_feature = torch.zeros([img_poly.size(0), cnn_feature.size(1), img_poly.size(1)]).to(img_poly.device)
+    for i in range(batch_size):
+        poly = img_poly[ind == i].unsqueeze(0)
+        gcn_feature[ind == i] = torch.nn.functional.grid_sample(cnn_feature[i:i + 1], poly)[0].permute(1, 0, 2)
+    return gcn_feature
+
+
+def get_adj_mat(n_adj, n_nodes):
+    a = np.zeros([n_nodes, n_nodes], dtype=np.float)
+
+    for i in range(n_nodes):
+        for j in range(-n_adj // 2, n_adj // 2 + 1):
+            if j != 0:
+                a[i][(i + j) % n_nodes] = 1
+                a[(i + j) % n_nodes][i] = 1
+    return a
+
+
+def get_adj_ind(n_adj, n_nodes, device):
+    ind = torch.tensor([i for i in range(-n_adj // 2, n_adj // 2 + 1) if i != 0]).long()
+    ind = (torch.arange(n_nodes)[:, None] + ind[None]) % n_nodes
+    return ind.to(device)
+
+
+def coord_embedding(b, w, h, device):
+    x_range = torch.linspace(0, 1, w, device=device)
+    y_range = torch.linspace(0, 1, h, device=device)
+    y, x = torch.meshgrid(y_range, x_range)
+    y = y.expand([b, 1, -1, -1])
+    x = x.expand([b, 1, -1, -1])
+    coord_map = torch.cat([x, y], 1)
+
+    return coord_map
+
+
+def img_poly_to_can_poly(img_poly):
+    if len(img_poly) == 0:
+        return torch.zeros_like(img_poly)
+    x_min = torch.min(img_poly[..., 0], dim=-1)[0]
+    y_min = torch.min(img_poly[..., 1], dim=-1)[0]
+    can_poly = img_poly.clone()
+    can_poly[..., 0] = can_poly[..., 0] - x_min[..., None]
+    can_poly[..., 1] = can_poly[..., 1] - y_min[..., None]
+    # x_max = torch.max(img_poly[..., 0], dim=-1)[0]
+    # y_max = torch.max(img_poly[..., 1], dim=-1)[0]
+    # h, w = y_max - y_min + 1, x_max - x_min + 1
+    # long_side = torch.max(h, w)
+    # can_poly = can_poly / long_side[..., None, None]
+    return can_poly

IndicPhotoOCR/detection/textbpn/network/layers/model_block.py

@@ -0,0 +1,149 @@
+# -*- coding: utf-8 -*-
+__author__ = "S.X.Zhang"
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from IndicPhotoOCR.detection.textbpn.network.layers.vgg import VggNet
+from IndicPhotoOCR.detection.textbpn.network.layers.resnet import ResNet
+from IndicPhotoOCR.detection.textbpn.network.layers.resnet_dcn import ResNet_DCN
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+
+
+class UpBlok(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.conv3x3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.deconv = nn.ConvTranspose2d(out_channels, out_channels, kernel_size=4, stride=2, padding=1)
+
+    def forward(self, upsampled, shortcut):
+        x = torch.cat([upsampled, shortcut], dim=1)
+        x = self.conv1x1(x)
+        x = F.relu(x)
+        x = self.conv3x3(x)
+        x = F.relu(x)
+        x = self.deconv(x)
+        return x
+
+
+class MergeBlok(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.conv3x3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, upsampled, shortcut):
+        x = torch.cat([upsampled, shortcut], dim=1)
+        x = self.conv1x1(x)
+        x = F.relu(x)
+        x = self.conv3x3(x)
+        return x
+
+
+class FPN(nn.Module):
+
+    def __init__(self, backbone='resnet50', is_training=True):
+        super().__init__()
+        self.is_training = is_training
+        self.backbone_name = backbone
+
+        if backbone in ['vgg_bn', 'vgg']:
+            self.backbone = VggNet(name=backbone, pretrain=is_training)
+            self.deconv5 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
+            self.merge4 = UpBlok(512 + 256, 128)
+            self.merge3 = UpBlok(256 + 128, 64)
+            if cfg.scale == 1:
+                self.merge2 = UpBlok(128 + 64, 32)  # FPN 1/2
+                self.merge1 = UpBlok(64 + 32, 32)   # FPN 1/1
+            elif cfg.scale == 2:
+                self.merge2 = UpBlok(128 + 64, 32)    # FPN 1/2
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/2
+            elif cfg.scale == 4:
+                self.merge2 = MergeBlok(128 + 64, 32)  # FPN 1/4
+
+        elif backbone in ['resnet50']:
+            self.backbone = ResNet(name=backbone, pretrain=is_training)
+            self.deconv5 = nn.ConvTranspose2d(2048, 256, kernel_size=4, stride=2, padding=1)
+            self.merge4 = UpBlok(1024 + 256, 128)
+            self.merge3 = UpBlok(512 + 128, 64)
+            if cfg.scale == 1:
+                self.merge2 = UpBlok(256 + 64, 32)  # FPN 1/2
+                self.merge1 = UpBlok(64 + 32, 32)   # FPN 1/1
+            elif cfg.scale == 2:
+                self.merge2 = UpBlok(256 + 64, 32)    # FPN 1/2
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/2
+            elif cfg.scale == 4:
+                self.merge2 = MergeBlok(256 + 64, 32)  # FPN 1/4
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/4
+        
+        elif backbone in ['resnet18']:
+            self.backbone = ResNet(name=backbone, pretrain=is_training)
+            self.deconv5 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
+            self.merge4 = UpBlok(256 + 256, 128)
+            self.merge3 = UpBlok(128 + 128, 64)
+            if cfg.scale == 1:
+                self.merge2 = UpBlok(64 + 64, 32)  # FPN 1/2
+                self.merge1 = UpBlok(64 + 32, 32)   # FPN 1/1
+            elif cfg.scale == 2:
+                self.merge2 = UpBlok(64 + 64, 32)    # FPN 1/2
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/2
+            elif cfg.scale == 4:
+                self.merge2 = MergeBlok(64 + 64, 32)  # FPN 1/4
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/4
+       
+        elif backbone in ["deformable_resnet18"]:
+            self.backbone = ResNet_DCN(name=backbone, pretrain=is_training)
+            self.deconv5 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
+            self.merge4 = UpBlok(256 + 256, 128)
+            self.merge3 = UpBlok(128 + 128, 64)
+            if cfg.scale == 1:
+                self.merge2 = UpBlok(64 + 64, 32)  # FPN 1/2
+                self.merge1 = UpBlok(64 + 32, 32)   # FPN 1/1
+            elif cfg.scale == 2:
+                self.merge2 = UpBlok(64 + 64, 32)    # FPN 1/2
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/2
+            elif cfg.scale == 4:
+                self.merge2 = MergeBlok(64 + 64, 32)  # FPN 1/4
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/4
+        
+        elif backbone in ["deformable_resnet50"]:
+            self.backbone = ResNet_DCN(name=backbone, pretrain=is_training)
+            self.deconv5 = nn.ConvTranspose2d(2048, 256, kernel_size=4, stride=2, padding=1)
+            self.merge4 = UpBlok(1024 + 256, 128)
+            self.merge3 = UpBlok(512 + 128, 64)
+            if cfg.scale == 1:
+                self.merge2 = UpBlok(256 + 64, 32)  # FPN 1/2
+                self.merge1 = UpBlok(64 + 32, 32)  # FPN 1/1
+            elif cfg.scale == 2:
+                self.merge2 = UpBlok(256 + 64, 32)  # FPN 1/2
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/2
+            elif cfg.scale == 4:
+                self.merge2 = MergeBlok(256 + 64, 32)  # FPN 1/4
+                self.merge1 = MergeBlok(64 + 32, 32)  # FPN 1/4
+        else:
+            print("backbone is not support !")
+
+    def forward(self, x):
+        C1, C2, C3, C4, C5 = self.backbone(x)
+        #print(C5.size())
+        #print(C4.size())
+        #print(C3.size())
+        #print(C2.size())
+        #print(C1.size())
+        up5 = self.deconv5(C5)
+        up5 = F.relu(up5)
+
+        up4 = self.merge4(C4, up5)
+        up4 = F.relu(up4)
+
+        up3 = self.merge3(C3, up4)
+        up3 = F.relu(up3)
+
+        up2 = self.merge2(C2, up3)
+        up2 = F.relu(up2)
+
+        up1 = self.merge1(C1, up2)
+        up1 = F.relu(up1)
+
+        return up1, up2, up3, up4, up5

IndicPhotoOCR/detection/textbpn/network/layers/position_encoding.py

@@ -0,0 +1,89 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+from util.misc import NestedTensor
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        mask = tensor_list.mask
+        assert mask is not None
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.row_embed = nn.Embedding(50, num_pos_feats)
+        self.col_embed = nn.Embedding(50, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        h, w = x.shape[-2:]
+        i = torch.arange(w, device=x.device)
+        j = torch.arange(h, device=x.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
+        return pos
+
+
+def build_position_encoding(args):
+    N_steps = args.hidden_dim // 2
+    if args.position_embedding in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    elif args.position_embedding in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(N_steps)
+    else:
+        raise ValueError(f"not supported {args.position_embedding}")
+
+    return position_embedding

IndicPhotoOCR/detection/textbpn/network/layers/resnet.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from torchvision.models import resnet
+import torch.utils.model_zoo as model_zoo
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+
+}
+
+
+class ResNet(nn.Module):
+    def __init__(self, name="resnet50", pretrain=True):
+        super().__init__()
+
+        if name == "resnet50":
+            base_net = resnet.resnet50(pretrained=False)
+        elif name == "resnet101":
+            base_net = resnet.resnet101(pretrained=False)
+        elif name == "resnet18":
+            base_net = resnet.resnet18(pretrained=False)
+        elif name == "resnet34":
+            base_net = resnet.resnet34(pretrained=False)
+
+        else:
+            print(" base model is not support !")
+
+        if pretrain:
+            print("load the {} weight from ./cache".format(name))
+            base_net.load_state_dict(model_zoo.load_url(model_urls[name], model_dir="./cache",
+                                                        map_location=torch.device(cfg.device)), strict=False)
+        # print(base_net)
+        self.stage1 = nn.Sequential(
+            base_net.conv1,
+            base_net.bn1,
+            base_net.relu,
+            base_net.maxpool
+        )
+        self.stage2 = base_net.layer1
+        self.stage3 = base_net.layer2
+        self.stage4 = base_net.layer3
+        self.stage5 = base_net.layer4
+        self.up2 = nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1)
+
+    def forward(self, x):
+        C1 = self.stage1(x)
+        C2 = self.stage2(C1)
+        C3 = self.stage3(C2)
+        C4 = self.stage4(C3)
+        C5 = self.stage5(C4)
+
+        if cfg.scale == 2 or cfg.scale == 1:
+            # up2 --> 1/2
+            C1 = self.up2(C1)
+
+        return C1, C2, C3, C4, C5
+
+
+if __name__ == '__main__':
+    import torch
+    input = torch.randn((4, 3, 512, 512))
+    net = ResNet()
+    C1, C2, C3, C4, C5 = net(input)
+    print(C1.size())
+    print(C2.size())
+    print(C3.size())
+    print(C4.size())
+    print(C5.size())

IndicPhotoOCR/detection/textbpn/network/layers/resnet_dcn.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from IndicPhotoOCR.detection.textbpn.network.backbone.resnet import deformable_resnet18,deformable_resnet50
+import torch.utils.model_zoo as model_zoo
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+
+}
+
+
+class ResNet_DCN(nn.Module):
+    def __init__(self, name="deformable_resnet18", pretrain=False):
+        super().__init__()
+
+        if name == "deformable_resnet18":
+            self.base_net = deformable_resnet18(pretrained=False)
+            if pretrain:
+                print("load the {} weight from ./cache".format(name))
+                self.base_net.load_state_dict(
+                    model_zoo.load_url(model_urls["resnet18"], model_dir="./cache",
+                                       map_location=torch.device(cfg.device)), strict=False)
+
+        elif name == "deformable_resnet50":
+            self.base_net = deformable_resnet50(pretrained=False)
+            if pretrain:
+                print("load the {} weight from ./cache".format(name))
+                self.base_net.load_state_dict(
+                    model_zoo.load_url(model_urls["resnet50"], model_dir="./cache",
+                                       map_location=torch.device(cfg.device)), strict=False)
+        else:
+            print(" base model is not support !")
+
+        # print(base_net)
+        self.up2 = nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1)
+
+    def forward(self, x):
+        C1, C2, C3, C4, C5 = self.base_net(x)
+        # up2 --> 1/2
+        C1 = self.up2(C1)
+
+        return C1, C2, C3, C4, C5
+
+
+if __name__ == '__main__':
+    import torch
+    input = torch.randn((4, 3, 512, 512))
+    net = ResNet_DCN()
+    C1, C2, C3, C4, C5 = net(input)
+    print(C1.size())
+    print(C2.size())
+    print(C3.size())
+    print(C4.size())
+    print(C5.size())

IndicPhotoOCR/detection/textbpn/network/layers/vgg.py

@@ -0,0 +1,62 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torchvision.models as models
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+
+model_urls = {
+    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
+    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
+    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
+    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
+    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
+    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
+}
+
+
+class VggNet(nn.Module):
+    def __init__(self, name="vgg16", pretrain=True):
+        super().__init__()
+        if name == "vgg16":
+            base_net = models.vgg16(pretrained=False)
+        elif name == "vgg16_bn":
+            base_net = models.vgg16_bn(pretrained=False)
+        else:
+            print(" base model is not support !")
+        if pretrain:
+            print("load the {} weight from ./cache".format(name))
+            base_net.load_state_dict(model_zoo.load_url(model_urls[name],
+                                                        model_dir="./cache",map_location=torch.device(cfg.device)))
+
+        if name == "vgg16":
+            self.stage1 = nn.Sequential(*[base_net.features[layer] for layer in range(0, 5)])
+            self.stage2 = nn.Sequential(*[base_net.features[layer] for layer in range(5, 10)])
+            self.stage3 = nn.Sequential(*[base_net.features[layer] for layer in range(10, 17)])
+            self.stage4 = nn.Sequential(*[base_net.features[layer] for layer in range(17, 24)])
+            self.stage5 = nn.Sequential(*[base_net.features[layer] for layer in range(24, 31)])
+        elif name == "vgg16_bn":
+            self.stage1 = nn.Sequential(*[base_net.features[layer] for layer in range(0, 7)])
+            self.stage2 = nn.Sequential(*[base_net.features[layer] for layer in range(7, 14)])
+            self.stage3 = nn.Sequential(*[base_net.features[layer] for layer in range(14, 24)])
+            self.stage4 = nn.Sequential(*[base_net.features[layer] for layer in range(24, 34)])
+            self.stage5 = nn.Sequential(*[base_net.features[layer] for layer in range(34, 44)])
+
+    def forward(self, x):
+        C1 = self.stage1(x)
+        C2 = self.stage2(C1)
+        C3 = self.stage3(C2)
+        C4 = self.stage4(C3)
+        C5 = self.stage5(C4)
+
+        return C1, C2, C3, C4, C5
+
+
+if __name__ == '__main__':
+    import torch
+    input = torch.randn((4, 3, 512, 512))
+    net = VggNet()
+    C1, C2, C3, C4, C5 = net(input)
+    print(C1.size())
+    print(C2.size())
+    print(C3.size())
+    print(C4.size())
+    print(C5.size())

IndicPhotoOCR/detection/textbpn/network/loss.py

@@ -0,0 +1,187 @@
+# -*- coding: utf-8 -*-
+# @Time    : 10/1/21
+# @Author  : GXYM
+import torch
+import torch.nn as nn
+from cfglib.config import config as cfg
+from network.Seg_loss import SegmentLoss
+from network.Reg_loss import PolyMatchingLoss
+import torch.nn.functional as F
+
+
+class TextLoss(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.MSE_loss = torch.nn.MSELoss(reduce=False, size_average=False)
+        self.BCE_loss = torch.nn.BCELoss(reduce=False, size_average=False)
+        self.PolyMatchingLoss = PolyMatchingLoss(cfg.num_points, cfg.device)
+        self.KL_loss = torch.nn.KLDivLoss(reduce=False, size_average=False)
+
+    @staticmethod
+    def single_image_loss(pre_loss, loss_label):
+        batch_size = pre_loss.shape[0]
+        sum_loss = torch.mean(pre_loss.view(-1)) * 0
+        pre_loss = pre_loss.view(batch_size, -1)
+        loss_label = loss_label.view(batch_size, -1)
+        eps = 0.001
+        for i in range(batch_size):
+            average_number = 0
+            positive_pixel = len(pre_loss[i][(loss_label[i] >= eps)])
+            average_number += positive_pixel
+            if positive_pixel != 0:
+                posi_loss = torch.mean(pre_loss[i][(loss_label[i] >= eps)])
+                sum_loss += posi_loss
+                if len(pre_loss[i][(loss_label[i] < eps)]) < 3 * positive_pixel:
+                    nega_loss = torch.mean(pre_loss[i][(loss_label[i] < eps)])
+                    average_number += len(pre_loss[i][(loss_label[i] < eps)])
+                else:
+                    nega_loss = torch.mean(torch.topk(pre_loss[i][(loss_label[i] < eps)], 3 * positive_pixel)[0])
+                    average_number += 3 * positive_pixel
+                sum_loss += nega_loss
+            else:
+                nega_loss = torch.mean(torch.topk(pre_loss[i], 100)[0])
+                average_number += 100
+                sum_loss += nega_loss
+            # sum_loss += loss/average_number
+
+        return sum_loss/batch_size
+
+    def cls_ohem(self, predict, target, train_mask, negative_ratio=3.):
+        pos = (target * train_mask).bool()
+        neg = ((1 - target) * train_mask).bool()
+
+        n_pos = pos.float().sum()
+        if n_pos.item() > 0:
+            loss_pos = self.BCE_loss(predict[pos], target[pos]).sum()
+            loss_neg = self.BCE_loss(predict[neg], target[neg])
+            n_neg = min(int(neg.float().sum().item()), int(negative_ratio * n_pos.float()))
+        else:
+            loss_pos = torch.tensor(0.)
+            loss_neg = self.BCE_loss(predict[neg], target[neg])
+            n_neg = 100
+        loss_neg, _ = torch.topk(loss_neg, n_neg)
+
+        return (loss_pos + loss_neg.sum()) / (n_pos + n_neg).float()
+
+    @staticmethod
+    def loss_calc_flux(pred_flux, gt_flux, weight_matrix, mask, train_mask):
+
+        # norm loss
+        gt_flux = 0.999999 * gt_flux / (gt_flux.norm(p=2, dim=1).unsqueeze(1) + 1e-3)
+        norm_loss = weight_matrix * torch.mean((pred_flux - gt_flux) ** 2, dim=1)*train_mask
+        norm_loss = norm_loss.sum(-1).mean()
+        # norm_loss = norm_loss.sum()
+
+        # angle loss
+        mask = train_mask * mask
+        pred_flux = 0.999999 * pred_flux / (pred_flux.norm(p=2, dim=1).unsqueeze(1) + 1e-3)
+        # angle_loss = weight_matrix * (torch.acos(torch.sum(pred_flux * gt_flux, dim=1))) ** 2
+        # angle_loss = angle_loss.sum(-1).mean()
+        angle_loss = (1 - torch.cosine_similarity(pred_flux, gt_flux, dim=1))
+        angle_loss = angle_loss[mask].mean()
+
+        return norm_loss, angle_loss
+
+    @staticmethod
+    def get_poly_energy(energy_field, img_poly, ind, h, w):
+        img_poly = img_poly.clone().float()
+        img_poly[..., 0] = img_poly[..., 0] / (w / 2.) - 1
+        img_poly[..., 1] = img_poly[..., 1] / (h / 2.) - 1
+
+        batch_size = energy_field.size(0)
+        gcn_feature = torch.zeros([img_poly.size(0), energy_field.size(1), img_poly.size(1)]).to(img_poly.device)
+        for i in range(batch_size):
+            poly = img_poly[ind == i].unsqueeze(0)
+            gcn_feature[ind == i] = torch.nn.functional.grid_sample(energy_field[i:i + 1], poly)[0].permute(1, 0, 2)
+        return gcn_feature
+
+    def loss_energy_regularization(self, energy_field, img_poly, inds, h, w):
+        energys = []
+        for i, py in enumerate(img_poly):
+            energy = self.get_poly_energy(energy_field.unsqueeze(1), py, inds, h, w)
+            energys.append(energy.squeeze(1).sum(-1))
+
+        regular_loss = torch.tensor(0.)
+        energy_loss = torch.tensor(0.)
+        for i, e in enumerate(energys[1:]):
+            regular_loss += torch.clamp(e - energys[i], min=0.0).mean()
+            energy_loss += torch.where(e <= 0.01, torch.tensor(0.), e).mean()
+
+        return (energy_loss+regular_loss)/len(energys[1:])
+
+    def forward(self, input_dict, output_dict, eps=None):
+        """
+          calculate boundary proposal network loss
+        """
+        # tr_mask = tr_mask.permute(0, 3, 1, 2).contiguous()
+
+        fy_preds = output_dict["fy_preds"]
+        py_preds = output_dict["py_preds"]
+        inds = output_dict["inds"]
+
+        train_mask = input_dict['train_mask']
+        tr_mask = input_dict['tr_mask'] > 0
+        distance_field = input_dict['distance_field']
+        direction_field = input_dict['direction_field']
+        weight_matrix = input_dict['weight_matrix']
+        gt_tags = input_dict['gt_points']
+
+        # # scale the prediction map
+        # fy_preds = F.interpolate(fy_preds, scale_factor=cfg.scale, mode='bilinear')
+        
+        if cfg.scale > 1:
+            train_mask = F.interpolate(train_mask.float().unsqueeze(1),
+                                       scale_factor=1/cfg.scale, mode='bilinear').squeeze().bool()
+            tr_mask = F.interpolate(tr_mask.float().unsqueeze(1),
+                                    scale_factor=1/cfg.scale, mode='bilinear').squeeze().bool()
+
+            distance_field = F.interpolate(distance_field.unsqueeze(1),
+                                           scale_factor=1/cfg.scale, mode='bilinear').squeeze()
+            direction_field = F.interpolate(direction_field,
+                                            scale_factor=1 / cfg.scale, mode='bilinear')
+            weight_matrix = F.interpolate(weight_matrix.unsqueeze(1),
+                                          scale_factor=1/cfg.scale, mode='bilinear').squeeze()
+
+        # pixel class loss
+        # cls_loss = self.cls_ohem(fy_preds[:, 0, :, :], tr_mask.float(), train_mask)
+        cls_loss = self.BCE_loss(fy_preds[:, 0, :, :],  tr_mask.float())
+        cls_loss = torch.mul(cls_loss, train_mask.float()).mean()
+
+        # distance field loss
+        dis_loss = self.MSE_loss(fy_preds[:, 1, :, :], distance_field)
+        dis_loss = torch.mul(dis_loss, train_mask.float())
+        dis_loss = self.single_image_loss(dis_loss, distance_field)
+
+        # # direction field loss
+        norm_loss, angle_loss = self.loss_calc_flux(fy_preds[:, 2:4, :, :], direction_field, 
+                                                           weight_matrix, tr_mask, train_mask)
+
+        # boundary point loss
+        point_loss = self.PolyMatchingLoss(py_preds[1:], gt_tags[inds])
+
+        #  Minimum energy loss regularization
+        h, w = distance_field.size(1) * cfg.scale, distance_field.size(2) * cfg.scale
+        energy_loss = self.loss_energy_regularization(distance_field, py_preds, inds[0], h, w)
+
+        if eps is None:
+            alpha = 1.0; beta = 3.0; theta=0.5; gama = 0.05
+        else:
+            alpha = 1.0; beta = 3.0; theta=0.5;
+            gama = 0.1*torch.sigmoid(torch.tensor((eps - cfg.max_epoch)/cfg.max_epoch))
+        loss = alpha*cls_loss + beta*dis_loss + theta*(norm_loss + angle_loss) + gama*(point_loss + energy_loss)
+
+        loss_dict = {
+            'total_loss': loss,
+            'cls_loss': alpha*cls_loss,
+            'distance loss': beta*dis_loss,
+            'dir_loss': theta*(norm_loss + angle_loss),
+            'norm_loss': theta*norm_loss,
+            'angle_loss': theta*angle_loss,
+            'point_loss': gama*point_loss,
+            'energy_loss': gama*energy_loss,
+
+        }
+
+        return loss_dict
+

IndicPhotoOCR/detection/textbpn/network/loss_org.py

@@ -0,0 +1,136 @@
+# -*- coding: utf-8 -*-
+# @Time    : 10/1/21
+# @Author  : GXYM
+import torch
+import torch.nn as nn
+from cfglib.config import config as cfg
+from network.Seg_loss import SegmentLoss
+from network.Reg_loss import PolyMatchingLoss
+
+
+class TextLoss(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.MSE_loss = torch.nn.MSELoss(reduce=False, size_average=False)
+        self.BCE_loss = torch.nn.BCELoss(reduce=False, size_average=False)
+        self.PolyMatchingLoss = PolyMatchingLoss(cfg.num_points, cfg.device)
+        self.KL_loss = torch.nn.KLDivLoss(reduce=False, size_average=False)
+
+    @staticmethod
+    def single_image_loss(pre_loss, loss_label):
+        batch_size = pre_loss.shape[0]
+        sum_loss = torch.mean(pre_loss.view(-1)) * 0
+        pre_loss = pre_loss.view(batch_size, -1)
+        loss_label = loss_label.view(batch_size, -1)
+        eps = 0.001
+        for i in range(batch_size):
+            average_number = 0
+            positive_pixel = len(pre_loss[i][(loss_label[i] >= eps)])
+            average_number += positive_pixel
+            if positive_pixel != 0:
+                posi_loss = torch.mean(pre_loss[i][(loss_label[i] >= eps)])
+                sum_loss += posi_loss
+                if len(pre_loss[i][(loss_label[i] < eps)]) < 3 * positive_pixel:
+                    nega_loss = torch.mean(pre_loss[i][(loss_label[i] < eps)])
+                    average_number += len(pre_loss[i][(loss_label[i] < eps)])
+                else:
+                    nega_loss = torch.mean(torch.topk(pre_loss[i][(loss_label[i] < eps)], 3 * positive_pixel)[0])
+                    average_number += 3 * positive_pixel
+                sum_loss += nega_loss
+            else:
+                nega_loss = torch.mean(torch.topk(pre_loss[i], 100)[0])
+                average_number += 100
+                sum_loss += nega_loss
+            # sum_loss += loss/average_number
+
+        return sum_loss/batch_size
+
+    def cls_ohem(self, predict, target, train_mask, negative_ratio=3.):
+        pos = (target * train_mask).bool()
+        neg = ((1 - target) * train_mask).bool()
+
+        n_pos = pos.float().sum()
+
+        if n_pos.item() > 0:
+            loss_pos = self.BCE_loss(predict[pos], target[pos]).sum()
+            loss_neg = self.BCE_loss(predict[neg], target[neg])
+            n_neg = min(int(neg.float().sum().item()), int(negative_ratio * n_pos.float()))
+        else:
+            loss_pos = torch.tensor(0.)
+            loss_neg = self.BCE_loss(predict[neg], target[neg])
+            n_neg = 100
+        loss_neg, _ = torch.topk(loss_neg, n_neg)
+
+        return (loss_pos + loss_neg.sum()) / (n_pos + n_neg).float()
+
+    @staticmethod
+    def loss_calc_flux(pred_flux, gt_flux, weight_matrix, mask, train_mask):
+
+        # norm loss
+        gt_flux = 0.999999 * gt_flux / (gt_flux.norm(p=2, dim=1).unsqueeze(1) + 1e-9)
+        norm_loss = weight_matrix * torch.sum((pred_flux - gt_flux) ** 2, dim=1)*train_mask
+        norm_loss = norm_loss.sum(-1).mean()
+
+        # angle loss
+        mask = train_mask * mask
+        pred_flux = 0.999999 * pred_flux / (pred_flux.norm(p=2, dim=1).unsqueeze(1) + 1e-9)
+        # angle_loss = weight_matrix * (torch.acos(torch.sum(pred_flux * gt_flux, dim=1))) ** 2
+        # angle_loss = angle_loss.sum(-1).mean()
+        angle_loss = (1 - torch.cosine_similarity(pred_flux, gt_flux, dim=1))
+        angle_loss = angle_loss[mask].mean()
+
+        return norm_loss, angle_loss
+
+    def forward(self, input_dict, output_dict, eps=None):
+        """
+          calculate boundary proposal network loss
+        """
+        # tr_mask = tr_mask.permute(0, 3, 1, 2).contiguous()
+
+        fy_preds = output_dict["fy_preds"]
+        py_preds = output_dict["py_preds"]
+        inds = output_dict["inds"]
+
+        train_mask = input_dict['train_mask']
+        tr_mask = input_dict['tr_mask'] > 0
+        distance_field = input_dict['distance_field']
+        direction_field = input_dict['direction_field']
+        weight_matrix = input_dict['weight_matrix']
+        gt_tags = input_dict['gt_points']
+
+        # pixel class loss
+        cls_loss = self.cls_ohem(fy_preds[:, 0, :, :], tr_mask.float(), train_mask.bool())
+
+        # distance field loss
+        dis_loss = self.MSE_loss(fy_preds[:, 1, :, :], distance_field)
+        dis_loss = torch.mul(dis_loss, train_mask.float())
+        dis_loss = self.single_image_loss(dis_loss, distance_field)
+
+        # direction field loss
+        norm_loss, angle_loss = self.loss_calc_flux(fy_preds[:, 2:4, :, :],
+                                                    direction_field, weight_matrix, tr_mask, train_mask)
+
+        # boundary point loss
+        point_loss = self.PolyMatchingLoss(py_preds, gt_tags[inds])
+
+        if eps is None:
+            loss_b = 0.05*point_loss
+            loss = cls_loss + 3.0*dis_loss + norm_loss + angle_loss + loss_b
+        else:
+            loss_b = 0.1*(torch.sigmoid(torch.tensor((eps - cfg.max_epoch)/cfg.max_epoch))) * point_loss
+            loss = cls_loss + 3.0*dis_loss + norm_loss + angle_loss + loss_b
+
+        loss_dict = {
+            'total_loss': loss,
+            'cls_loss': cls_loss,
+            'distance loss': 3.0*dis_loss,
+            'dir_loss': norm_loss + angle_loss,
+            'point_loss': loss_b,
+            'norm_loss': norm_loss,
+            'angle_loss': angle_loss,
+
+        }
+
+        return loss_dict
+

IndicPhotoOCR/detection/textbpn/network/textnet.py

@@ -0,0 +1,216 @@
+# -*- coding: utf-8 -*-
+# @Time    : 10/1/21
+# @Author  : GXYM
+import torch
+import torch.nn as nn
+from IndicPhotoOCR.detection.textbpn.network.layers.model_block import FPN
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+import numpy as np
+from IndicPhotoOCR.detection.textbpn.network.layers.CircConv import DeepSnake
+from IndicPhotoOCR.detection.textbpn.network.layers.GCN import GCN
+from IndicPhotoOCR.detection.textbpn.network.layers.RNN import RNN
+from IndicPhotoOCR.detection.textbpn.network.layers.Adaptive_Deformation import AdaptiveDeformation
+# from IndicPhotoOCR.detection.textbpn.network.layers.Transformer_old import Transformer_old
+from IndicPhotoOCR.detection.textbpn.network.layers.Transformer import Transformer
+import cv2
+from IndicPhotoOCR.detection.textbpn.util.misc import get_sample_point, fill_hole
+from IndicPhotoOCR.detection.textbpn.network.layers.gcn_utils import get_node_feature, \
+    get_adj_mat, get_adj_ind, coord_embedding, normalize_adj
+import torch.nn.functional as F
+import time
+
+
+class Evolution(nn.Module):
+    def __init__(self, node_num, adj_num, is_training=True, device=None, model="snake"):
+        super(Evolution, self).__init__()
+        self.node_num = node_num
+        self.adj_num = adj_num
+        self.device = device
+        self.is_training = is_training
+        self.clip_dis = 16
+
+        self.iter = 3
+        if model == "gcn":
+            self.adj = get_adj_mat(self.adj_num, self.node_num)
+            self.adj = normalize_adj(self.adj, type="DAD").float().to(self.device)
+            for i in range(self.iter):
+                evolve_gcn = GCN(36, 128)
+                self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+        elif model == "rnn":
+            self.adj = None
+            for i in range(self.iter):
+                evolve_gcn = RNN(36, 128)
+                self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+        elif model == "AD":
+            self.adj = get_adj_mat(self.adj_num, self.node_num)
+            self.adj = normalize_adj(self.adj, type="DAD").float().to(self.device)
+            for i in range(self.iter):
+                evolve_gcn = AdaptiveDeformation(36, 128)
+                self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+        # elif model == "BT_old":
+        #     self.adj = None
+        #     for i in range(self.iter):
+        #         evolve_gcn = Transformer_old(36, 512, num_heads=8,
+        #                                  dim_feedforward=2048, drop_rate=0.0, if_resi=True, block_nums=4)
+        #         self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+        elif model == "BT":
+            self.adj = None
+            for i in range(self.iter):
+                evolve_gcn = Transformer(36, 128, num_heads=8,
+                                         dim_feedforward=1024, drop_rate=0.0, if_resi=True, block_nums=3)
+                self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+        else:
+            self.adj = get_adj_ind(self.adj_num, self.node_num, self.device)
+            for i in range(self.iter):
+                evolve_gcn = DeepSnake(state_dim=128, feature_dim=36, conv_type='dgrid')
+                self.__setattr__('evolve_gcn' + str(i), evolve_gcn)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d):
+                m.weight.data.normal_(0.0, 0.02)
+                # nn.init.kaiming_normal_(m.weight, mode='fan_in')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    @staticmethod
+    def get_boundary_proposal(input=None, seg_preds=None, switch="gt"):
+
+        if switch == "gt":
+            inds = torch.where(input['ignore_tags'] > 0)
+            # if len(inds[0]) > 320:
+            #    inds = (inds[0][:320], inds[1][:320])
+            init_polys = input['proposal_points'][inds]
+        else:
+            tr_masks = input['tr_mask'].cpu().numpy()
+            tcl_masks = seg_preds[:, 0, :, :].detach().cpu().numpy() > cfg.threshold
+            inds = []
+            init_polys = []
+            for bid, tcl_mask in enumerate(tcl_masks):
+                ret, labels = cv2.connectedComponents(tcl_mask.astype(np.uint8), connectivity=8)
+                for idx in range(1, ret):
+                    text_mask = labels == idx
+                    ist_id = int(np.sum(text_mask*tr_masks[bid])/np.sum(text_mask))-1
+                    inds.append([bid, ist_id])
+                    poly = get_sample_point(text_mask, cfg.num_points, cfg.approx_factor)
+                    init_polys.append(poly)
+            inds = torch.from_numpy(np.array(inds)).permute(1, 0).to(input["img"].device)
+            init_polys = torch.from_numpy(np.array(init_polys)).to(input["img"].device)
+
+        return init_polys, inds, None
+
+    def get_boundary_proposal_eval(self, input=None, seg_preds=None):
+
+        # if cfg.scale > 1:
+        #     seg_preds = F.interpolate(seg_preds, scale_factor=cfg.scale, mode='bilinear')
+        cls_preds = seg_preds[:, 0, :, :].detach().cpu().numpy()
+        dis_preds = seg_preds[:, 1, :, ].detach().cpu().numpy()
+
+        inds = []
+        init_polys = []
+        confidences = []
+        for bid, dis_pred in enumerate(dis_preds):
+            # # dis_mask = (dis_pred / np.max(dis_pred)) > cfg.dis_threshold
+            dis_mask = dis_pred > cfg.dis_threshold
+            # dis_mask = fill_hole(dis_mask)
+            ret, labels = cv2.connectedComponents(dis_mask.astype(np.uint8), connectivity=8, ltype=cv2.CV_16U)
+            for idx in range(1, ret):
+                text_mask = labels == idx
+                confidence = round(cls_preds[bid][text_mask].mean(), 3)
+                # 50 for MLT2017 and ArT (or DCN is used in backone); else is all 150;
+                # just can set to 50, which has little effect on the performance
+                if np.sum(text_mask) < 50/(cfg.scale*cfg.scale) or confidence < cfg.cls_threshold:
+                    continue
+                confidences.append(confidence)
+                inds.append([bid, 0])
+                
+                poly = get_sample_point(text_mask, cfg.num_points,
+                                        cfg.approx_factor, scales=np.array([cfg.scale, cfg.scale]))
+                init_polys.append(poly)
+
+        if len(inds) > 0:
+            inds = torch.from_numpy(np.array(inds)).permute(1, 0).to(input["img"].device, non_blocking=True)
+            init_polys = torch.from_numpy(np.array(init_polys)).to(input["img"].device, non_blocking=True).float()
+        else:
+            init_polys = torch.from_numpy(np.array(init_polys)).to(input["img"].device, non_blocking=True).float()
+            inds = torch.from_numpy(np.array(inds)).to(input["img"].device, non_blocking=True)
+
+        return init_polys, inds, confidences
+
+    def evolve_poly(self, snake, cnn_feature, i_it_poly, ind):
+        if len(i_it_poly) == 0:
+            return torch.zeros_like(i_it_poly)
+        h, w = cnn_feature.size(2)*cfg.scale, cnn_feature.size(3)*cfg.scale
+        node_feats = get_node_feature(cnn_feature, i_it_poly, ind, h, w)
+        i_poly = i_it_poly + torch.clamp(snake(node_feats, self.adj).permute(0, 2, 1), -self.clip_dis, self.clip_dis)
+        if self.is_training:
+            i_poly = torch.clamp(i_poly, 0, w-1)
+        else:
+            i_poly[:, :, 0] = torch.clamp(i_poly[:, :, 0], 0, w - 1)
+            i_poly[:, :, 1] = torch.clamp(i_poly[:, :, 1], 0, h - 1)
+        return i_poly
+
+    def forward(self, embed_feature, input=None, seg_preds=None, switch="gt"):
+        if self.is_training:
+            init_polys, inds, confidences = self.get_boundary_proposal(input=input, seg_preds=seg_preds, switch=switch)
+            # TODO sample fix number
+        else:
+            init_polys, inds, confidences = self.get_boundary_proposal_eval(input=input, seg_preds=seg_preds)
+            if init_polys.shape[0] == 0:
+                return [init_polys for i in range(self.iter+1)], inds, confidences
+
+        py_preds = [init_polys, ]
+        for i in range(self.iter):
+            evolve_gcn = self.__getattr__('evolve_gcn' + str(i))
+            init_polys = self.evolve_poly(evolve_gcn, embed_feature, init_polys, inds[0])
+            py_preds.append(init_polys)
+
+        return py_preds, inds, confidences
+
+
+class TextNet(nn.Module):
+
+    def __init__(self, backbone='vgg', is_training=True):
+        super().__init__()
+        self.is_training = is_training
+        self.backbone_name = backbone
+        self.fpn = FPN(self.backbone_name, is_training=(not cfg.resume and is_training))
+
+        self.seg_head = nn.Sequential(
+            nn.Conv2d(32, 16, kernel_size=3, padding=2, dilation=2),
+            nn.PReLU(),
+            nn.Conv2d(16, 16, kernel_size=3, padding=4, dilation=4),
+            nn.PReLU(),
+            nn.Conv2d(16, 4, kernel_size=1, stride=1, padding=0),
+        )
+        self.BPN = Evolution(cfg.num_points, adj_num=4,
+                             is_training=is_training, device=cfg.device, model="BT")
+
+    def load_model(self, model_path):
+        print('Loading from {}'.format(model_path))
+        state_dict = torch.load(model_path, map_location=torch.device(cfg.device))
+        self.load_state_dict(state_dict['model'], strict=(not self.is_training))
+
+    def forward(self, input_dict, test_speed=False):
+        output = {}
+        b, c, h, w = input_dict["img"].shape
+        if self.is_training or cfg.exp_name in ['ArT', 'MLT2017', "MLT2019"] or test_speed:
+            image = input_dict["img"]
+        else:
+            image = torch.zeros((b, c, cfg.test_size[1], cfg.test_size[1]), dtype=torch.float32).to(cfg.device)
+            image[:, :, :h, :w] = input_dict["img"][:, :, :, :]
+
+        up1, _, _, _, _ = self.fpn(image)
+        up1 = up1[:, :, :h // cfg.scale, :w // cfg.scale]
+
+        preds = self.seg_head(up1)
+        fy_preds = torch.cat([torch.sigmoid(preds[:, 0:2, :, :]), preds[:, 2:4, :, :]], dim=1)
+        cnn_feats = torch.cat([up1, fy_preds], dim=1)
+
+        py_preds, inds, confidences = self.BPN(cnn_feats, input=input_dict, seg_preds=fy_preds, switch="gt")
+        
+        output["fy_preds"] = fy_preds
+        output["py_preds"] = py_preds
+        output["inds"] = inds
+        output["confidences"] = confidences
+
+        return output

IndicPhotoOCR/detection/textbpn/textbpnpp_detector.py

@@ -0,0 +1,197 @@
+import torch
+import cv2
+import numpy as np
+from IndicPhotoOCR.detection.textbpn.network.textnet import TextNet
+from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
+import warnings
+import os
+import requests
+from tqdm import tqdm
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+model_info = {
+    "textbpnpp": {
+        "path": "models/TextBPN_resnet50_300.pth",
+        "url" : "https://github.com/Bhashini-IITJ/SceneTextDetection/releases/download/TextBPN%2B%2B/TextBPN_resnet50_300.pth",
+    },
+    "textbpnpp_deformable": {
+        "path":"models/TextBPN_deformable_resnet50_300.pth",
+        "url": "https://github.com/Bhashini-IITJ/SceneTextDetection/releases/download/TextBPN%2B%2B/TextBPN_deformable_resnet50_300.pth",
+    },
+    "textbpn_resnet18" : {
+        "path":"models/TextBPN_resnet18_300.pth",
+        "url": "https://github.com/Bhashini-IITJ/SceneTextDetection/releases/download/TextBPN%2B%2B/TextBPN_resnet18_300.pth",
+
+    }
+}
+        # Ensure model file exists; download directly if not
+def ensure_model(model_name):
+    model_path = model_info[model_name]["path"]
+    url = model_info[model_name]["url"]
+    root_model_dir = "IndicPhotoOCR/detection/textbpn"
+    model_path = os.path.join(root_model_dir, model_path)
+    
+    if not os.path.exists(model_path):
+        print(f"Model not found locally. Downloading {model_name} from {url}...")
+        
+        # Start the download with a progress bar
+        response = requests.get(url, stream=True)
+        total_size = int(response.headers.get('content-length', 0))
+        os.makedirs(f"{root_model_dir}/models", exist_ok=True)
+        
+        with open(model_path, "wb") as f, tqdm(
+                desc=model_name,
+                total=total_size,
+                unit='B',
+                unit_scale=True,
+                unit_divisor=1024,
+        ) as bar:
+            for data in response.iter_content(chunk_size=1024):
+                f.write(data)
+                bar.update(len(data))
+
+        print(f"Downloaded model for {model_name}.")
+        
+    return model_path
+
+class TextBPNpp_detector:
+    def __init__(self, model_name="textbpnpp", backbone="resnet50", device="cpu"):
+        """
+        Initialize the TextBPN model.
+        :param model_path: Path to the pre-trained model.
+        :param backbone: Backbone architecture (default: "resnet50").
+        :param device: Device to run the model on (default: "cpu").
+        """
+        self.model_path = ensure_model(model_name)
+        self.device = torch.device(device)
+        self.model = TextNet(is_training=False, backbone=backbone)
+        self.model.load_model(self.model_path)
+        self.model.eval()
+        self.model.to(self.device)
+
+    @staticmethod
+    def to_device(tensor, device):
+        """
+        Move tensor to the specified device.
+        :param tensor: Tensor to move.
+        :param device: Target device.
+        :return: Tensor on the target device.
+        """
+        return tensor.to(device, non_blocking=True)
+
+    @staticmethod
+    def pad_image(image, stride=32):
+        """
+        Pad the image to make its dimensions divisible by the stride.
+        :param image: Input image.
+        :param stride: Stride size.
+        :return: Padded image and original dimensions.
+        """
+        h, w = image.shape[:2]
+        new_h = (h + stride - 1) // stride * stride
+        new_w = (w + stride - 1) // stride * stride
+        padded_image = cv2.copyMakeBorder(
+            image, 0, new_h - h, 0, new_w - w, cv2.BORDER_CONSTANT, value=(0, 0, 0)
+        )
+        return padded_image, (h, w)
+
+    @staticmethod
+    def rescale_result(image, bbox_contours, original_height, original_width):
+        """
+        Rescale the bounding box contours to the original image size.
+        :param image: Image after resizing.
+        :param bbox_contours: Bounding box contours.
+        :param original_height: Original image height.
+        :param original_width: Original image width.
+        :return: Original image and rescaled contours.
+        """
+        contours = []
+        for cont in bbox_contours:
+            cont[:, 0] = (cont[:, 0] * original_width / image.shape[1]).astype(int)
+            cont[:, 1] = (cont[:, 1] * original_height / image.shape[0]).astype(int)
+            contours.append(cont)
+        return contours
+
+    def detect(self, image_path):
+        """
+        Perform text detection on the given image.
+        :param image_path: Path to the input image.
+        :return: Dictionary with detection results.
+        """
+        image = cv2.imread(image_path)
+        if image is None:
+            raise ValueError(f"Failed to read the image at {image_path}")
+
+        padded_image, original_size = self.pad_image(image)
+        padded_tensor = (
+            torch.from_numpy(padded_image).permute(2, 0, 1).float() / 255.0
+        ).unsqueeze(0)  # Convert to tensor and add batch dimension
+
+        cfg.test_size = [padded_image.shape[0], padded_image.shape[1]]
+
+        input_dict = {"img": self.to_device(padded_tensor, self.device)}
+        with torch.no_grad():
+            output_dict = self.model(input_dict, padded_image.shape)
+
+        contours = output_dict["py_preds"][-1].int().cpu().numpy()
+        contours = self.rescale_result(image, contours, *original_size)
+
+        bbox_result_dict = {"detections": []}
+        for contour in contours:
+            # x_min, y_min = np.min(contour, axis=0)
+            # x_max, y_max = np.max(contour, axis=0)
+            # bbox_result_dict["detections"].append([x_min, y_min, x_max, y_max])
+            bbox_result_dict["detections"].append(contour.tolist())
+
+        return bbox_result_dict
+
+    def visualize_detections(self, image_path, bbox_result_dict, output_path="output.png"):
+        """
+        Visualize detections on the image.
+        :param image_path: Path to the input image.
+        :param bbox_result_dict: Detection results in the format:
+                                {'detections': [[[x1, y1], [x2, y2], [x3, y3], [x4, y4]], ...]}.
+        :param output_path: Path to save the visualized image. If None, the image is only displayed.
+        """
+        # Load the image
+        image = cv2.imread(image_path)
+        if image is None:
+            raise ValueError(f"Failed to read the image at {image_path}")
+        
+        # Draw each detection
+        for bbox in bbox_result_dict.get("detections", []):
+            points = np.array(bbox, dtype=np.int32)  # Convert to numpy array
+            cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)
+        
+        # Display or save the visualized image
+        if output_path:
+            cv2.imwrite(output_path, image)
+            print(f"Visualization saved to {output_path}")
+        else:
+            cv2.imshow("Detections", image)
+            cv2.waitKey(0)
+            cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description='Text detection using EAST model')
+    parser.add_argument('--image_path', type=str, required=True, help='Path to the input image')
+    parser.add_argument('--device', type=str, default='cpu', help='Device to run the model on, e.g., "cpu" or "cuda"')
+    parser.add_argument('--model_name', type=str, required=True, help='Path to the model checkpoint file')
+    args = parser.parse_args()
+
+
+
+    # model_path = "/DATA1/ocrteam/anik/git/IndicPhotoOCR/IndicPhotoOCR/detection/textbpn/models/TextBPN_resnet50_300.pth"
+    # image_path = "/DATA1/ocrteam/anik/splitonBSTD/detection/D/image_542.jpg"
+
+    detector = TextBPNpp_detector(args.model_name, device="cpu")
+    result = detector.detect(args.image_path)
+    print(result)
+    # detector.visualize_detections(image_path, result)
+
+    # python -m IndicPhotoOCR.detection.textbpn.textbpnpp_detector \
+    # --image_path /DATA1/ocrteam/anik/splitonBSTD/detection/D/image_542.jpg \
+    # --model_name textbpnpp

IndicPhotoOCR/detection/textbpn/util/__init__.py

@@ -0,0 +1,2 @@
+from .visualize import *
+from .pbox import *

IndicPhotoOCR/detection/textbpn/util/augmentation.py

@@ -0,0 +1,794 @@
+# -*- coding: utf-8 -*-
+__author__ = "S.X.Zhang"
+import numpy as np
+import math
+import cv2
+import copy
+import numpy.random as random
+from shapely.geometry import Polygon
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as F
+from PIL import ImageEnhance, Image
+
+
+###<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<###
+###<<<<<<<<<  Function  >>>>>>>>>>>>###
+###>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>###
+def crop_first(image, polygons, scale =10):
+    polygons_new = copy.deepcopy(polygons)
+    h, w, _ = image.shape
+    pad_h = h // scale
+    pad_w = w // scale
+    h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
+    w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
+
+    text_polys = []
+    pos_polys = []
+    for polygon in polygons_new:
+        rect = cv2.minAreaRect(polygon.points.astype(np.int32))
+        box = cv2.boxPoints(rect)
+        box = np.int0(box)
+        text_polys.append([box[0], box[1], box[2], box[3]])
+        if polygon.label != -1:
+            pos_polys.append([box[0], box[1], box[2], box[3]])
+
+    polys = np.array(text_polys, dtype=np.int32)
+    for poly in polys:
+        poly = np.round(poly, decimals=0).astype(np.int32)  # 四舍五入
+        minx = np.min(poly[:, 0])
+        maxx = np.max(poly[:, 0])
+        w_array[minx + pad_w:maxx + pad_w] = 1
+        miny = np.min(poly[:, 1])
+        maxy = np.max(poly[:, 1])
+        h_array[miny + pad_h:maxy + pad_h] = 1
+    # ensure the cropped area not across a text 保证截取区域不会横穿文字
+    h_axis = np.where(h_array == 0)[0]
+    w_axis = np.where(w_array == 0)[0]
+    pp_polys = np.array(pos_polys, dtype=np.int32)
+
+    return h_axis, w_axis, pp_polys
+
+####<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<####
+####<<<<<<<<<<<  Class  >>>>>>>>>>>>>####
+####>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>####
+class Compose(object):
+    """Composes several augmentations together.
+    Args:
+        transforms (List[Transform]): list of transforms to compose.
+    Example:
+        >>> augmentations.Compose([
+        >>>     transforms.CenterCrop(10),
+        >>>     transforms.ToTensor(),
+        >>> ])
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img, pts=None):
+        for t in self.transforms:
+            img, pts = t(img, pts)
+        return img, pts
+
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = np.array(mean)
+        self.std = np.array(std)
+
+    def __call__(self, image, polygons=None):
+        image = image.astype(np.float32)
+        image /= 255.0
+        image -= self.mean
+        image /= self.std
+        return image, polygons
+
+
+class MinusMean(object):
+    def __init__(self, mean):
+        self.mean = np.array(mean)
+
+    def __call__(self, image, polygons=None):
+        image = image.astype(np.float32)
+        image -= self.mean
+        return image, polygons
+
+
+class RandomMirror(object):
+    # 镜像
+    def __init__(self):
+        pass
+
+    def __call__(self, image, polygons=None):
+        if polygons is None:
+            return image, polygons
+        if random.random()< 0.3:
+            image = np.ascontiguousarray(image[:, ::-1])
+            _, width, _ = image.shape
+            for polygon in polygons:
+                polygon.points[:, 0] = width - polygon.points[:, 0]
+        return image, polygons
+
+
+class AugmentColor(object):
+    # 颜色增强（添加噪声）
+    def __init__(self):
+        self.U = np.array([[-0.56543481, 0.71983482, 0.40240142],
+                      [-0.5989477, -0.02304967, -0.80036049],
+                      [-0.56694071, -0.6935729, 0.44423429]], dtype=np.float32)
+        self.EV = np.array([1.65513492, 0.48450358, 0.1565086], dtype=np.float32)
+        self.sigma = 0.1
+        self.color_vec = None
+
+    def __call__(self, img, polygons=None):
+        color_vec = self.color_vec
+        if self.color_vec is None:
+            if not self.sigma > 0.0:
+                color_vec = np.zeros(3, dtype=np.float32)
+            else:
+                color_vec = np.random.normal(0.0, self.sigma, 3)
+
+        alpha = color_vec.astype(np.float32) * self.EV
+        noise = np.dot(self.U, alpha.T) * 255
+        return np.clip(img + noise[np.newaxis, np.newaxis, :], 0, 255), polygons
+
+
+class RandomContrast(object):
+    def __init__(self, lower=0.5, upper=1.5):
+        self.lower = lower
+        self.upper = upper
+        assert self.upper >= self.lower, "contrast upper must be >= lower."
+        assert self.lower >= 0, "contrast lower must be non-negative."
+
+    # expects float image
+    def __call__(self, image, polygons=None):
+        if random.randint(2):
+            alpha = random.uniform(self.lower, self.upper)
+            image *= alpha
+        return np.clip(image, 0, 255), polygons
+
+
+class RandomBrightness(object):
+    def __init__(self, delta=32):
+        assert delta >= 0.0
+        assert delta <= 255.0
+        self.delta = delta
+
+    def __call__(self, image, polygons=None):
+        image = image.astype(np.float32)
+        if random.randint(2):
+            delta = random.uniform(-self.delta, self.delta)
+            image += delta
+        return np.clip(image, 0, 255), polygons
+
+
+class RandomErasing(object):
+    def __init__(self, sr=(0.0004, 0.01), scale=(0.5, 3), ratio=0.2, Type ="Erasing"):
+        """
+
+        :param area:
+        :param type: Erasing or Cutout
+        """
+        self.sr = sr
+        self.scale= scale
+        self.ratio=ratio
+        self.type=Type
+
+    def __call__(self, img, polygons=None):
+
+        if random.random()< self.ratio:
+            return img, polygons
+        area=img.shape[0]*img.shape[1]
+        target_area=random.randint(*self.sr)*area
+        aspect_ratio=random.uniform(*self.scale)
+        h = int(round(math.sqrt(target_area / aspect_ratio)))
+        w = int(round(math.sqrt(target_area * aspect_ratio)))
+
+        if w < img.shape[1] and h < img.shape[0]:
+            x1 = random.randint(0, img.shape[1] - w)
+            y1 = random.randint(0, img.shape[0] - h)
+            if self.type == "Erasing":
+                color=(random.randint(0, 255),random.randint(0, 255),random.randint(0, 255))
+                img[y1:y1+h, x1:x1+h,:]=color
+            else:
+                Gray_value=random.randint(0, 255)
+                color = (Gray_value, Gray_value ,Gray_value)
+                img[y1:y1 + h, x1:x1 + h, :] = color
+
+        return img, polygons
+
+
+class RandomMixUp(object):
+    def __init__(self, mixup_alpha=2):
+        self.mixup_alpha = mixup_alpha
+
+    def __call__(self, img1, img2, label1=[], label2=[]):
+        beta=np.random.beta(self.mixup_alpha,self.mixup_alpha)
+
+        #image = img1 * Gama + (1 - Gama) * img2
+        image=cv2.addWeighted(img1, beta, img2, (1-beta), 0)
+
+        if label1 is None or label2  is None:
+            return  img1, label1
+        if isinstance(label1, list) and isinstance(label2, list):
+            label=[]
+            for id in range(len(label1)):
+                lab = beta*label1[id]+ (1-beta)*label2[id]
+                label.append(lab)
+            return image, label
+        else:
+            print("Error: label is not a list type")
+
+        return img1, label1
+
+
+class Rotate(object):
+    def __init__(self, up=30):
+        self.up = up
+
+    @staticmethod
+    def rotate(center, pt, theta):  # 二维图形学的旋转
+        xr, yr = center
+        yr = -yr
+        x, y = pt[:, 0], pt[:, 1]
+        y = -y
+
+        theta = theta / 180 * math.pi
+        cos = math.cos(theta)
+        sin = math.sin(theta)
+
+        _x = xr + (x - xr) * cos - (y - yr) * sin
+        _y = yr + (x - xr) * sin + (y - yr) * cos
+
+        return _x, -_y
+
+    def __call__(self, img, polygons=None):
+        if np.random.randint(2):
+            return img, polygons
+        angle = np.random.normal(loc=0.0, scale=0.5) * self.up  # angle 按照高斯分布
+        rows, cols = img.shape[0:2]
+        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1.0)
+        img = cv2.warpAffine(img, M, (cols, rows), borderValue=[0, 0, 0])
+        center = cols / 2.0, rows / 2.0
+        if polygons is not None:
+            for polygon in polygons:
+                x, y = self.rotate(center, polygon.points, angle)
+                pts = np.vstack([x, y]).T
+                polygon.points = pts
+        return img, polygons
+
+
+class RotatePadding(object):
+    def __init__(self, up=60,colors=True):
+        self.up = up
+        self.colors = colors
+        self.ratio = 0.5
+
+    @staticmethod
+    def rotate(center, pt, theta, movSize=[0, 0], scale=1):  # 二维图形学的旋转
+        (xr, yr) = center
+        yr = -yr
+        x, y = pt[:, 0], pt[:, 1]
+        y = -y
+
+        theta = theta / 180 * math.pi
+        cos = math.cos(theta)
+        sin = math.sin(theta)
+
+        x = (x - xr) * scale
+        y = (y - yr) * scale
+
+        _x = xr + x * cos - y * sin + movSize[0]
+        _y = -(yr + x * sin + y * cos) + movSize[1]
+
+        return _x, _y
+
+    @staticmethod
+    def shift(size, degree):
+        angle = degree * math.pi / 180.0
+        width = size[0]
+        height = size[1]
+
+        alpha = math.cos(angle)
+        beta = math.sin(angle)
+        new_width = int(width * math.fabs(alpha) + height * math.fabs(beta))
+        new_height = int(width * math.fabs(beta) + height * math.fabs(alpha))
+
+        size = [new_width, new_height]
+        return size
+
+    def __call__(self, image, polygons=None, scale=1.0):
+        if np.random.random() <= self.ratio:
+            return image, polygons
+        angle = np.random.normal(loc=0.0, scale=0.5) * self.up  # angle 按照高斯分布
+        rows, cols = image.shape[0:2]
+        center = (cols / 2.0, rows / 2.0)
+        newSize = self.shift([cols * scale, rows * scale], angle)
+        movSize = [int((newSize[0] - cols) / 2), int((newSize[1] - rows) / 2)]
+
+        M = cv2.getRotationMatrix2D(center, angle, scale)
+        M[0, 2] += int((newSize[0] - cols) / 2)
+        M[1, 2] += int((newSize[1] - rows) / 2)
+
+        if self.colors:
+            H, W, _ = image.shape
+            mask = np.zeros_like(image)
+            (h_index, w_index) = (np.random.randint(0, H * 7 // 8), np.random.randint(0, W * 7 // 8))
+            img_cut = image[h_index:(h_index + H // 9), w_index:(w_index + W // 9)]
+            img_cut = cv2.resize(img_cut, (newSize[0], newSize[1]))
+            mask = cv2.warpAffine(mask, M, (newSize[0], newSize[1]), borderValue=[1, 1, 1])
+            image = cv2.warpAffine(image, M, (newSize[0], newSize[1]), borderValue=[0,0,0])
+            image=image+img_cut*mask
+        else:
+            color = [0, 0, 0]
+            image = cv2.warpAffine(image, M, (newSize[0], newSize[1]), borderValue=color)
+
+        if polygons is not None:
+            for polygon in polygons:
+                x, y = self.rotate(center, polygon.points, angle,movSize,scale)
+                pts = np.vstack([x, y]).T
+                polygon.points = pts
+        return image, polygons
+
+
+class SquarePadding(object):
+
+    def __call__(self, image, polygons=None):
+
+        H, W, _ = image.shape
+
+        if H == W:
+            return image, polygons
+
+        padding_size = max(H, W)
+        (h_index, w_index) = (np.random.randint(0, H*7//8),np.random.randint(0, W*7//8))
+        img_cut = image[h_index:(h_index+H//9),w_index:(w_index+W//9)]
+        expand_image = cv2.resize(img_cut,(padding_size, padding_size))
+        #expand_image = np.zeros((padding_size, padding_size, 3), dtype=image.dtype)
+        #expand_image=img_cut[:,:,:]
+        if H > W:
+            y0, x0 = 0, (H - W) // 2
+        else:
+            y0, x0 = (W - H) // 2, 0
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points += np.array([x0, y0])
+        expand_image[y0:y0+H, x0:x0+W] = image
+        image = expand_image
+
+        return image, polygons
+
+
+class RandomImgCropPatch(object):
+    def __init__(self, up=30, beta=0.3):
+        self.up = up
+        self.beta=0.3
+        self.scale = 10
+
+    @staticmethod
+    def get_contour_min_area_box(contour):
+        rect = cv2.minAreaRect(contour)
+        box = cv2.boxPoints(rect)
+        box = np.int0(box)
+        return box
+
+    def CropWH(self, image, cut_w, cut_h, polygons=None):
+        h_axis, w_axis, polys = crop_first(image, polygons, scale=self.scale)
+        h, w, _ = image.shape
+        pad_h = h // self.scale
+        pad_w = w // self.scale
+        # TODO try Flip
+        xx = np.random.choice(w_axis, size=2)
+        xmin = np.min(xx) - pad_w
+        xmax = xmin + cut_w
+        yy = np.random.choice(h_axis, size=2)
+        ymin = np.min(yy) - pad_h
+        ymax = ymin + cut_h
+        if polys.shape[0] != 0:
+            poly_axis_in_area = (polys[:, :, 0] >= xmin) & (polys[:, :, 0] <= xmax) \
+                                & (polys[:, :, 1] >= ymin) & (polys[:, :, 1] <= ymax)
+            selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
+        else:
+            selected_polys = []
+
+        cropped = image[ymin:ymax + 1, xmin:xmax + 1, :]
+        polygons_new = []
+        for idx in selected_polys:
+            polygon = polygons[idx]
+            polygon.points -= np.array([xmin, ymin])
+            polygons_new.append(polygon)
+        image = cropped
+        polygon = polygons_new
+
+        return image, polygon
+
+    def __call__(self, images, polygons_list=None):
+        I_x, I_y = 1024,1024
+
+        w = int(round(I_x * random.beta(self.beta, self.beta)))
+        h = int(round(I_y * random.beta(self.beta, self.beta)))
+        w_ = [w, I_x - w, w, I_x - w]
+        h_ = [h, h, I_y - h, I_y - h]
+        new_img = np.zeros((I_x, I_y, 3), dtype=images[0].dtype)
+        imgs=[]
+        new_polygons=[]
+        for i, im in enumerate(images):
+           img, polygons = self.CropWH(im,  w_[i],  h_[i], polygons=polygons_list[i])
+           imgs.append(img)
+           new_polygons.append(polygons)
+        new_img[0:w, 0:h, :] = imgs[0]
+        new_img[w:I_x, 0:h, :] = imgs[1]
+        new_img[0:w, h:I_y, :] = imgs[2]
+        new_img[w:I_x, h:I_y, :] = imgs[3]
+        for polygon in new_polygons[1]:
+            polygon.points += np.array([w, 0])
+        for polygon in new_polygons[2]:
+            polygon.points += np.array([0, h])
+        for polygon in new_polygons[3]:
+            polygon.points += np.array([w, h])
+
+        polygons=new_polygons[0]+new_polygons[1]+new_polygons[2]+new_polygons[3]
+
+        return new_img, polygons
+
+
+class RandomCropFlip(object):
+
+    def __init__(self, min_crop_side_ratio=0.01):
+        self.scale = 10
+        self.ratio = 0.2
+        self.epsilon = 10.0
+        self.min_crop_side_ratio = min_crop_side_ratio
+
+    def __call__(self, image, polygons=None):
+
+        if polygons is None:
+            return image, polygons
+
+        if np.random.random() <= self.ratio:
+            return image, polygons
+
+        # 计算 有效的Crop区域, 方便选取有效的种子点
+        h_axis, w_axis, pp_polys = crop_first(image, polygons, scale =self.scale)
+        if len(h_axis) == 0 or len(w_axis) == 0:
+            return image, polygons
+
+        # TODO try crop
+        attempt = 0
+        h, w, _ = image.shape
+        area = h * w
+        pad_h = h // self.scale
+        pad_w = w // self.scale
+        while attempt < 10:
+            attempt += 1
+            polygons_new = []
+            xx = np.random.choice(w_axis, size=2)
+            xmin = np.min(xx) - pad_w
+            xmax = np.max(xx) - pad_w
+            xmin = np.clip(xmin, 0, w - 1)
+            xmax = np.clip(xmax, 0, w - 1)
+            yy = np.random.choice(h_axis, size=2)
+            ymin = np.min(yy) - pad_h
+            ymax = np.max(yy) - pad_h
+            ymin = np.clip(ymin, 0, h - 1)
+            ymax = np.clip(ymax, 0, h - 1)
+            if (xmax - xmin) * (ymax - ymin) < area * self.min_crop_side_ratio:
+                # area too small
+                continue
+
+            pts = np.stack([[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]).T.astype(np.int32)
+            pp = Polygon(pts).buffer(0)
+            Fail_flag = False
+            for polygon in polygons:
+                ppi = Polygon(polygon.points).buffer(0)
+                ppiou = float(ppi.intersection(pp).area)
+                if np.abs(ppiou - float(ppi.area)) > self.epsilon and np.abs(ppiou) > self.epsilon:
+                    Fail_flag = True
+                    break
+                if np.abs(ppiou - float(ppi.area)) < self.epsilon:
+                    polygons_new.append(polygon)
+
+            if Fail_flag:
+                continue
+            else:
+                break
+
+        if len(polygons_new) == 0:
+            cropped = image[ymin:ymax, xmin:xmax, :]
+            select_type = random.randint(3)
+            if select_type == 0:
+                img = np.ascontiguousarray(cropped[:, ::-1])
+            elif select_type == 1:
+                img = np.ascontiguousarray(cropped[::-1, :])
+            else:
+                img = np.ascontiguousarray(cropped[::-1, ::-1])
+            image[ymin:ymax, xmin:xmax, :] = img
+            return image, polygons
+
+        else:
+            cropped = image[ymin:ymax, xmin:xmax, :]
+            height, width, _ = cropped.shape
+            select_type = random.randint(3)
+            if select_type == 0:
+                img = np.ascontiguousarray(cropped[:, ::-1])
+                for polygon in polygons_new:
+                    polygon.points[:, 0] = width - polygon.points[:, 0] + 2 * xmin
+            elif select_type == 1:
+                img = np.ascontiguousarray(cropped[::-1, :])
+                for polygon in polygons_new:
+                    polygon.points[:, 1] = height - polygon.points[:, 1] + 2 * ymin
+            else:
+                img = np.ascontiguousarray(cropped[::-1, ::-1])
+                for polygon in polygons_new:
+                    polygon.points[:, 0] = width - polygon.points[:, 0] + 2 * xmin
+                    polygon.points[:, 1] = height - polygon.points[:, 1] + 2 * ymin
+            image[ymin:ymax, xmin:xmax, :] = img
+
+        return image, polygons
+
+
+class RandomResizedCrop(object):
+    def __init__(self, min_crop_side_ratio=0.1):
+        self.scale = 10
+        self.epsilon = 1e-2
+        self.min_crop_side_ratio = min_crop_side_ratio
+
+    def __call__(self, image, polygons):
+
+        if polygons is None:
+            return image, polygons
+
+        # 计算 有效的Crop区域, 方便选取有效的种子点
+        h_axis, w_axis, pp_polys = crop_first(image, polygons, scale =self.scale)
+        if len(h_axis) == 0 or len(w_axis) == 0:
+            return image, polygons
+
+        # TODO try crop
+        attempt = 0
+        h, w, _ = image.shape
+        area = h * w
+        pad_h = h // self.scale
+        pad_w = w // self.scale
+        while attempt < 10:
+            attempt += 1
+            xx = np.random.choice(w_axis, size=2)
+            xmin = np.min(xx) - pad_w
+            xmax = np.max(xx) - pad_w
+            xmin = np.clip(xmin, 0, w - 1)
+            xmax = np.clip(xmax, 0, w - 1)
+            yy = np.random.choice(h_axis, size=2)
+            ymin = np.min(yy) - pad_h
+            ymax = np.max(yy) - pad_h
+            ymin = np.clip(ymin, 0, h - 1)
+            ymax = np.clip(ymax, 0, h - 1)
+            if (xmax - xmin)*(ymax - ymin) <area*self.min_crop_side_ratio:
+                # area too small
+                continue
+            if pp_polys.shape[0] != 0:
+                poly_axis_in_area = (pp_polys[:, :, 0] >= xmin) & (pp_polys[:, :, 0] <= xmax) \
+                                    & (pp_polys[:, :, 1] >= ymin) & (pp_polys[:, :, 1] <= ymax)
+                selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
+            else:
+                selected_polys = []
+
+            if len(selected_polys) == 0:
+                continue
+            else:
+                pts = np.stack([[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]).T.astype(np.int32)
+                pp = Polygon(pts).buffer(0)
+                polygons_new = []
+                Fail_flag = False
+                for polygon in copy.deepcopy(polygons):
+                    ppi = Polygon(polygon.points).buffer(0)
+                    ppiou = float(ppi.intersection(pp).area)
+                    if np.abs(ppiou - float(ppi.area)) > self.epsilon and np.abs(ppiou) > self.epsilon:
+                        Fail_flag = True
+                        break
+                    elif np.abs(ppiou - float(ppi.area)) < self.epsilon:
+                        # polygon.points -= np.array([xmin, ymin])
+                        polygons_new.append(polygon)
+
+                if Fail_flag:
+                    continue
+                else:
+                    cropped = image[ymin:ymax + 1, xmin:xmax + 1, :]
+                    for polygon in polygons_new:
+                        polygon.points -= np.array([xmin, ymin])
+
+                    return cropped, polygons_new
+
+        return image, polygons
+
+
+class RandomResizeScale(object):
+    def __init__(self, size=512, ratio=(3./4, 5./2)):
+        self.size = size
+        self.ratio = ratio
+
+    def __call__(self, image, polygons=None):
+
+        aspect_ratio = np.random.uniform(self.ratio[0], self.ratio[1])
+        h, w, _ = image.shape
+        scales = self.size*1.0/max(h, w)
+        aspect_ratio = scales * aspect_ratio
+        aspect_ratio = int(w * aspect_ratio)*1.0/w
+        image = cv2.resize(image, (int(w * aspect_ratio), int(h*aspect_ratio)))
+        scales = np.array([aspect_ratio, aspect_ratio])
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points = polygon.points * scales
+
+        return image, polygons
+
+
+class Resize(object):
+    def __init__(self, size=1024):
+        self.size = size
+        self.SP = SquarePadding()
+
+    def __call__(self, image, polygons=None):
+        h, w, _ = image.shape
+        image = cv2.resize(image, (self.size,
+                                   self.size))
+        scales = np.array([self.size / w, self.size / h])
+
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points = polygon.points * scales
+
+        return image, polygons
+
+
+class ResizeSquare(object):
+    def __init__(self, size=(480, 1280)):
+        self.size = size
+
+    def __call__(self, image, polygons=None):
+        h, w, _ = image.shape
+        img_size_min = min(h, w)
+        img_size_max = max(h, w)
+
+        if img_size_min < self.size[0]:
+            im_scale = float(self.size[0]) / float(img_size_min)  # expand min to size[0]
+            if np.ceil(im_scale * img_size_max) > self.size[1]:  # expand max can't > size[1]
+                im_scale = float(self.size[1]) / float(img_size_max)
+        elif img_size_max > self.size[1]:
+            im_scale = float(self.size[1]) / float(img_size_max)
+        else:
+            im_scale = 1.0
+
+        new_h = int(int(h * im_scale/32)*32)
+        new_w = int(int(w * im_scale/32)*32)
+        # if new_h*new_w > 1600*1920:
+        #     im_scale = 1600 / float(img_size_max)
+        #     new_h = int(int(h * im_scale/32)*32)
+        #     new_w = int(int(w * im_scale/32)*32)
+        image = cv2.resize(image, (new_w, new_h))
+        scales = np.array([new_w / w, new_h / h])
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points = polygon.points * scales
+
+        return image, polygons
+
+
+class ResizeLimitSquare(object):
+    def __init__(self, size=512, ratio=0.6):
+        self.size = size
+        self.ratio = ratio
+        self.SP = SquarePadding()
+
+    def __call__(self, image, polygons=None):
+        if np.random.random() <= self.ratio:
+            image, polygons = self.SP(image, polygons)
+        h, w, _ = image.shape
+        image = cv2.resize(image, (self.size,self.size))
+        scales = np.array([self.size*1.0/ w, self.size*1.0 / h])
+
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points = polygon.points * scales
+
+        return image, polygons
+
+
+class RandomResizePadding(object):
+    def __init__(self, size=512, random_scale=np.array([0.75, 1.0, 1.25,1.5,2.0]),stride=32, ratio=0.6667):
+        self.random_scale = random_scale
+        self.size = size
+        self.ratio=ratio
+        self.stride=stride
+        self.SP=SquarePadding()
+
+        ###########Random size for different eproches ########################
+        rd_scale = np.random.choice(self.random_scale)
+        step_num = round(np.random.normal(loc=0.0, scale=0.35) * 8)  # step 按照高斯分布
+        self.input_size = np.clip(int(self.size * rd_scale + step_num * self.stride),
+                                  (int(self.size * self.random_scale[0] - self.stride)),
+                                  int(self.size * self.random_scale[-1] + self.stride))
+        ############################ end ########################
+
+    def __call__(self, image, polygons=None):
+
+        if np.random.random() <= self.ratio:
+            image, polygons = self.SP(image, polygons)
+        h, w, _ = image.shape
+        image = cv2.resize(image, (self.input_size,self.input_size))
+        scales = np.array([self.input_size*1.0/ w, self.input_size*1.0 / h])
+
+        if polygons is not None:
+            for polygon in polygons:
+                polygon.points = polygon.points * scales
+
+        return image, polygons
+
+transform_type_dict = dict(
+    brightness=ImageEnhance.Brightness, contrast=ImageEnhance.Contrast,
+    sharpness=ImageEnhance.Sharpness,   color=ImageEnhance.Color
+)
+
+
+class RandomDistortion(object):
+    def __init__(self, transform_dict, prob=0.5):
+        self.transforms = [(transform_type_dict[k], transform_dict[k]) for k in transform_dict]
+        self.prob = prob
+
+    def __call__(self, img, target):
+        if random.random() > self.prob:
+            return img, target
+        out = Image.fromarray(img)
+        rand_num = np.random.uniform(0, 1, len(self.transforms))
+
+        for i, (transformer, alpha) in enumerate(self.transforms):
+            r = alpha * (rand_num[i] * 2.0 - 1.0) + 1  # r in [1-alpha, 1+alpha)
+            out = transformer(out).enhance(r)
+
+        return np.array(out), target
+
+
+class Augmentation(object):
+    def __init__(self, size, mean, std):
+        self.size = size
+        self.mean = mean
+        self.std = std
+        self._transform_dict = {'brightness': 0.5, 'contrast': 0.5, 'sharpness': 0.8386, 'color': 0.5}
+        self.augmentation = Compose([
+            RandomCropFlip(),
+            RandomResizeScale(size=self.size, ratio=(3. / 8, 5. / 2)),
+            RandomResizedCrop(),
+            RotatePadding(up=60, colors=True),  # pretrain on Syn is "up=30", else is "up=60"
+            ResizeLimitSquare(size=self.size),
+            RandomMirror(),
+            RandomDistortion(self._transform_dict),
+            Normalize(mean=self.mean, std=self.std),
+        ])
+
+    def __call__(self, image, polygons=None):
+        return self.augmentation(image, polygons)
+
+
+class BaseTransform(object):
+    def __init__(self, size, mean, std):
+        self.size = size
+        self.mean = mean
+        self.std = std
+        self.augmentation = Compose([
+            # Resize(size=640),
+            ResizeSquare(size=self.size),
+            Normalize(mean, std)
+        ])
+
+    def __call__(self, image, polygons=None):
+        return self.augmentation(image, polygons)
+
+
+class BaseTransformNresize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+        self.augmentation = Compose([
+            Normalize(mean, std)
+        ])
+
+    def __call__(self, image, polygons=None):
+        return self.augmentation(image, polygons)

IndicPhotoOCR/detection/textbpn/util/canvas.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+__author__ = '古溪'
+
+import numpy as np
+import random
+import matplotlib.pyplot as plt
+
+
+def heatmap(im_gray):
+    cmap = plt.get_cmap('jet')
+    rgba_img = cmap(255 - im_gray)
+    Hmap = np.delete(rgba_img, 3, 2)
+    # print(Hmap.shape, Hmap.max(), Hmap.min())
+    # cv2.imshow("heat_img", Hmap)
+    # cv2.waitKey(0)
+    return Hmap
+
+
+def loss_ploy(loss_list, steps, period, name=""):
+    fig1, ax1 = plt.subplots(figsize=(16, 9))
+    ax1.plot(range(steps // period), loss_list)
+    ax1.set_title("Average loss vs step*{}".format(period))
+    ax1.set_xlabel("step*{}".format(period))
+    ax1.set_ylabel("Current loss")
+    plt.savefig('{}@loss_vs_step*{}.png'.format(name,period))
+    plt.clf()
+
+
+def plt_ploys(ploys, period, name=""):
+    fig1, ax1 = plt.subplots(figsize=(16, 9))
+    cnames = ['aliceblue','antiquewhite','aqua','aquamarine','azure',
+               'blanchedalmond','blue','blueviolet','brown','burlywood',
+               'coral','cornflowerblue','cornsilk','crimson','cyan',
+               'darkblue','deeppink','deepskyblue','dodgerblue','forestgreen',
+               'gold','goldenrod','green','greenyellow','honeydew','hotpink',
+               'lawngreen','lightblue','lightgreen','lightpink','lightsalmon',
+               'lightseagreen','lightsteelblue','lightyellow','lime','limegreen',
+               'mediumseagreen','mediumspringgreen','midnightblue','orange','orangered',
+               'pink','red','royalblue','seagreen','skyblue','springgreen','steelblue',
+               'tan','teal','thistle','yellow','yellowgreen']
+
+    color = random.sample(cnames, len(ploys.keys()))
+    for ii, key in enumerate(ploys.keys()):
+        ax1.plot(range(1, len(ploys[key])+1), ploys[key],color=color[ii], label=key)
+    ax1.set_title("Loss Carve line")
+    ax1.set_xlabel("step*{}".format(period))
+    ax1.set_ylabel("Current loss")
+    plt.legend(ploys.keys())
+    plt.savefig('{}@loss_vs_step*{}.png'.format(name, period))
+    plt.clf()
+
+if __name__ == '__main__':
+    # TODO ADD CODE
+    pass

IndicPhotoOCR/detection/textbpn/util/detection.py

@@ -0,0 +1,48 @@
+# c++ version pse based on opencv 3+
+from pse import decode as pse_decode
+from cfglib.config import config as cfg
+
+
+class TextDetector(object):
+
+    def __init__(self, model):
+        # evaluation mode
+        self.model = model
+        model.eval()
+        # parameter
+        self.scale = cfg.scale
+        self.threshold = cfg.threshold
+
+    def detect(self, image, img_show):
+        # get model output
+        preds = self.model.forward(image)
+        preds, boxes, contours = pse_decode(preds[0], self.scale, self.threshold)
+
+        output = {
+            'image': image,
+            'tr': preds,
+            'bbox': boxes
+        }
+        return contours, output
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

IndicPhotoOCR/detection/textbpn/util/eval.py

@@ -0,0 +1,228 @@
+import os
+import cv2
+import numpy as np
+import subprocess
+from cfglib.config import config as cfg
+from util.misc import mkdirs
+
+
+def osmkdir(out_dir):
+    import shutil
+    if os.path.exists(out_dir):
+        shutil.rmtree(out_dir)
+    os.makedirs(out_dir)
+
+
+def analysize_result(source_dir, fid_path, outpt_dir, name):
+
+    bad_txt = open("{}/eval.txt".format(outpt_dir), 'w')
+    all_eval = open("{}/{}/{}_eval.txt".format(cfg.output_dir, "Analysis", name), 'a+')
+    sel_list = list()
+    with open(fid_path) as f:
+        lines = f.read().split("\n")
+        for line in lines:
+            line_items = line.split(" ")
+            id = line_items[0]
+            precision = float(line_items[2].split('=')[-1])
+            recall = float(line_items[4].split('=')[-1])
+            if id != "ALL" and (precision < 0.5 or recall < 0.5):
+                img_path = os.path.join(source_dir, line_items[0].replace(".txt", ".jpg"))
+                if os.path.exists(img_path):
+                    os.system('cp {} {}'.format(img_path, outpt_dir))
+                sel_list.append((int(id.replace(".txt", "").replace("img", "").replace("_", "")), line))
+            if id == "ALL":
+                all_eval.write("{} {} {}\n".format(
+                    outpt_dir.split('/')[-1],
+                    "{}/{}".format(cfg.dis_threshold, cfg.cls_threshold),
+                    line))
+    sel_list = sorted(sel_list, key=lambda its: its[0])
+    bad_txt.write('\n'.join([its[1] for its in sel_list]))
+    all_eval.close()
+    bad_txt.close()
+
+
+def deal_eval_total_text(debug=False):
+    # compute DetEval
+    eval_dir = os.path.join(cfg.output_dir, "Analysis", "output_eval")
+    if not os.path.exists(eval_dir):
+        os.makedirs(eval_dir)
+
+    print('Computing DetEval in {}/{}'.format(cfg.output_dir, cfg.exp_name))
+    subprocess.call(
+        ['python', 'dataset/total_text/Evaluation_Protocol/Python_scripts/Deteval.py', cfg.exp_name, '--tr', '0.7',
+         '--tp', '0.6'])
+    subprocess.call(
+        ['python', 'dataset/total_text/Evaluation_Protocol/Python_scripts/Deteval.py', cfg.exp_name, '--tr', '0.8',
+         '--tp', '0.4'])
+
+    if debug:
+        source_dir = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name))
+        outpt_dir_base = os.path.join(cfg.output_dir, "Analysis", "eval_view", "total_text")
+        if not os.path.exists(outpt_dir_base):
+            mkdirs(outpt_dir_base)
+
+        outpt_dir1 = os.path.join(outpt_dir_base, "{}_{}_{}_{}_{}"
+                                  .format(cfg.test_size[0], cfg.test_size[1], cfg.checkepoch, 0.7, 0.6))
+        osmkdir(outpt_dir1)
+        fid_path1 = '{}/Eval_TotalText_{}_{}.txt'.format(eval_dir, 0.7, 0.6)
+
+        analysize_result(source_dir, fid_path1, outpt_dir1, "totalText")
+
+        outpt_dir2 = os.path.join(outpt_dir_base, "{}_{}_{}_{}_{}"
+                                  .format(cfg.test_size[0], cfg.test_size[1], cfg.checkepoch, 0.8, 0.4))
+        osmkdir(outpt_dir2)
+        fid_path2 = '{}/Eval_TotalText_{}_{}.txt'.format(eval_dir, 0.8, 0.4)
+
+        analysize_result(source_dir, fid_path2, outpt_dir2, "totalText")
+
+    print('End.')
+
+
+def deal_eval_ctw1500(debug=False):
+    # compute DetEval
+    eval_dir = os.path.join(cfg.output_dir, "Analysis", "output_eval")
+    if not os.path.exists(eval_dir):
+        os.makedirs(eval_dir)
+
+    print('Computing DetEval in {}/{}'.format(cfg.output_dir, cfg.exp_name))
+    subprocess.call(['python', 'dataset/ctw1500/Evaluation_Protocol/ctw1500_eval.py', cfg.exp_name])
+
+    if debug:
+        source_dir = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name))
+        outpt_dir_base = os.path.join(cfg.output_dir, "Analysis", "eval_view", "ctw1500")
+        if not os.path.exists(outpt_dir_base):
+            mkdirs(outpt_dir_base)
+
+        outpt_dir = os.path.join(outpt_dir_base, "{}_{}_{}".format(cfg.test_size[0], cfg.test_size[1], cfg.checkepoch))
+        osmkdir(outpt_dir)
+        fid_path1 = '{}/Eval_ctw1500_{}.txt'.format(eval_dir, 0.5)
+
+        analysize_result(source_dir, fid_path1, outpt_dir, "ctw1500")
+
+    print('End.')
+
+
+def deal_eval_icdar15(debug=False):
+    # compute DetEval
+    eval_dir = os.path.join(cfg.output_dir, "Analysis", "output_eval")
+    if not os.path.exists(eval_dir):
+        os.makedirs(eval_dir)
+
+    input_dir = 'output/{}'.format(cfg.exp_name)
+    father_path = os.path.abspath(input_dir)
+    print(father_path)
+    print('Computing DetEval in {}/{}'.format(cfg.output_dir, cfg.exp_name))
+    subprocess.call(['sh', 'dataset/icdar15/eval.sh', father_path])
+
+    if debug:
+        source_dir = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name))
+        outpt_dir_base = os.path.join(cfg.output_dir, "Analysis", "eval_view", "icdar15")
+        if not os.path.exists(outpt_dir_base):
+            mkdirs(outpt_dir_base)
+
+        outpt_dir = os.path.join(outpt_dir_base, "{}_{}_{}".format(cfg.test_size[0], cfg.test_size[1], cfg.checkepoch))
+        osmkdir(outpt_dir)
+        fid_path1 = '{}/Eval_icdar15.txt'.format(eval_dir)
+
+        analysize_result(source_dir, fid_path1, outpt_dir, "icdar15")
+
+    print('End.')
+
+    pass
+
+
+def deal_eval_TD500(debug=False):
+    # compute DetEval
+    eval_dir = os.path.join(cfg.output_dir, "Analysis", "output_eval")
+    if not os.path.exists(eval_dir):
+        os.makedirs(eval_dir)
+
+    input_dir = 'output/{}'.format(cfg.exp_name)
+    father_path = os.path.abspath(input_dir)
+    print(father_path)
+    print('Computing DetEval in {}/{}'.format(cfg.output_dir, cfg.exp_name))
+    subprocess.call(['sh', 'dataset/TD500/eval.sh', father_path])
+
+    if debug:
+        source_dir = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name))
+        outpt_dir_base = os.path.join(cfg.output_dir, "Analysis", "eval_view", "TD500")
+        if not os.path.exists(outpt_dir_base):
+            mkdirs(outpt_dir_base)
+
+        outpt_dir = os.path.join(outpt_dir_base, "{}_{}_{}".format(cfg.test_size[0], cfg.test_size[1], cfg.checkepoch))
+        osmkdir(outpt_dir)
+        fid_path1 = '{}/Eval_TD500.txt'.format(eval_dir)
+
+        analysize_result(source_dir, fid_path1, outpt_dir, "TD500")
+
+    print('End.')
+
+
+def data_transfer_ICDAR(contours):
+    cnts = list()
+    for cont in contours:
+        rect = cv2.minAreaRect(cont)
+        if min(rect[1][0], rect[1][1]) <= 5:
+            continue
+        points = cv2.boxPoints(rect)
+        points = np.int0(points)
+        # print(points.shape)
+        # points = np.reshape(points, (4, 2))
+        cnts.append(points)
+    return cnts
+
+
+def data_transfer_TD500(contours, res_file, img=None):
+    with open(res_file, 'w') as f:
+        for cont in contours:
+            rect = cv2.minAreaRect(cont)
+            if min(rect[1][0], rect[1][1]) <= 5:
+                continue
+            points = cv2.boxPoints(rect)
+            box = np.int0(points)
+            cv2.drawContours(img, [box], 0, (0, 255, 0), 3)
+
+            cx, cy = rect[0]
+            w_, h_ = rect[1]
+            angle = rect[2]
+            mid_ = 0
+            if angle > 45:
+                angle = 90 - angle
+                mid_ = w_;
+                w_ = h_;
+                h_ = mid_
+            elif angle < -45:
+                angle = 90 + angle
+                mid_ = w_;
+                w_ = h_;
+                h_ = mid_
+            angle = angle / 180 * 3.141592653589
+
+            x_min = int(cx - w_ / 2)
+            x_max = int(cx + w_ / 2)
+            y_min = int(cy - h_ / 2)
+            y_max = int(cy + h_ / 2)
+            f.write('{},{},{},{},{}\r\n'.format(x_min, y_min, x_max, y_max, angle))
+
+    return img
+
+
+def data_transfer_MLT2017(contours, res_file):
+    with open(res_file, 'w') as f:
+        for cont in contours:
+            rect = cv2.minAreaRect(cont)
+            if min(rect[1][0], rect[1][1]) <= 5:
+                 continue
+            ploy_area = cv2.contourArea(cont)
+            rect_area = rect[1][0]*rect[1][1]
+            solidity = ploy_area/rect_area
+            width = rect[1][0] - np.clip(rect[1][0] * (1-np.sqrt(solidity)), 0, 6)
+            height = rect[1][1] - np.clip(rect[1][1] * (1-np.sqrt(solidity)), 0, 4)
+            points = cv2.boxPoints((rect[0], (width, height), rect[2]))
+            points = np.int0(points)
+            p = np.reshape(points, -1)
+            f.write('{},{},{},{},{},{},{},{},{}\r\n'
+                    .format(p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], 1))
+
+
+

IndicPhotoOCR/detection/textbpn/util/graph.py

@@ -0,0 +1,309 @@
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+
+import numpy as np
+import time
+from util.misc import norm2
+
+class Data(object):
+    def __init__(self, name):
+        self.__name = name
+        self.__links = set()
+
+    @property
+    def name(self):
+        return self.__name
+
+    @property
+    def links(self):
+        return set(self.__links)
+
+    def add_link(self, other, score):
+        self.__links.add(other)
+        other.__links.add(self)
+
+
+def connected_components(nodes, score_dict, th):
+    '''
+    conventional connected components searching
+    '''
+    result = []
+    nodes = set(nodes)
+    while nodes:
+        n = nodes.pop()
+        group = {n}
+        queue = [n]
+        while queue:
+            n = queue.pop(0)
+            if th is not None:
+                neighbors = {l for l in n.links if score_dict[tuple(sorted([n.name, l.name]))] >= th}
+            else:
+                neighbors = n.links
+            neighbors.difference_update(group)
+            nodes.difference_update(neighbors)
+            group.update(neighbors)
+            queue.extend(neighbors)
+        result.append(group)
+    return result
+
+
+def connected_components_constraint(nodes, max_sz, score_dict=None, th=None):
+    '''
+    only use edges whose scores are above `th`
+    if a component is larger than `max_sz`, all the nodes in this component are added into `remain` and returned for next iteration.
+    '''
+    result = []
+    remain = set()
+    nodes = set(nodes)
+    while nodes:
+        n = nodes.pop()
+        group = {n}
+        queue = [n]
+        valid = True
+        while queue:
+            n = queue.pop(0)
+            if th is not None:
+                neighbors = {l for l in n.links if score_dict[tuple(sorted([n.name, l.name]))] >= th}
+            else:
+                neighbors = n.links
+            neighbors.difference_update(group)
+            nodes.difference_update(neighbors)
+            group.update(neighbors)
+            queue.extend(neighbors)
+            if len(group) > max_sz or len(remain.intersection(neighbors)) > 0:
+                # if this group is larger than `max_sz`, add the nodes into `remain`
+                valid = False
+                remain.update(group)
+                break
+        if valid: # if this group is smaller than or equal to `max_sz`, finalize it.
+            result.append(group)
+    return result, remain
+
+
+def graph_propagation_naive(edges, score, th, bboxs=None, dis_thresh=50, pool='avg'):
+
+    edges = np.sort(edges, axis=1)
+
+    score_dict = {}  # score lookup table
+    if pool is None:
+        for i, e in enumerate(edges):
+            score_dict[e[0], e[1]] = score[i]
+    elif pool == 'avg':
+        for i, e in enumerate(edges):
+            if bboxs is not None:
+                box1 = bboxs[e[0]][:8].reshape(4, 2)
+                box2 = bboxs[e[1]][:8].reshape(4, 2)
+                c1 = np.mean(box1, 0); c2 = np.mean(box2, 0)
+                dst = norm2(c1 - c2)
+                if dst > dis_thresh:
+                    score[i] = 0
+            if (e[0], e[1]) in score_dict:
+                score_dict[e[0], e[1]] = 0.5 * (score_dict[e[0], e[1]] + score[i])
+            else:
+                score_dict[e[0], e[1]] = score[i]
+
+    elif pool == 'max':
+        for i, e in enumerate(edges):
+            if (e[0], e[1]) in score_dict:
+                score_dict[e[0], e[1]] = max(score_dict[e[0], e[1]], score[i])
+            else:
+                score_dict[e[0], e[1]] = score[i]
+    else:
+        raise ValueError('Pooling operation not supported')
+
+    nodes = np.sort(np.unique(edges.flatten()))
+    mapping = -1 * np.ones((nodes.max()+1), dtype=np.int)
+    mapping[nodes] = np.arange(nodes.shape[0])
+    link_idx = mapping[edges]
+    vertex = [Data(n) for n in nodes]
+    for l, s in zip(link_idx, score):
+        vertex[l[0]].add_link(vertex[l[1]], s)
+
+    # first iteration
+    comps = connected_components(vertex, score_dict,th)
+
+    return comps
+
+
+def graph_search(edges, scores, edges_num, th=None):
+    # graph search
+    scores = scores.reshape((-1, edges_num))
+    select_index = np.argsort(scores, axis=1)[:, -2:]
+    edges = np.sort(edges, axis=1).reshape((-1, edges_num, 2))
+
+    score_dict = {}
+    for i, ips in enumerate(select_index):
+        edg = edges[i]
+        si = scores[i]
+        for j, idx in enumerate(ips):
+            e = edg[idx, :]
+            if (e[0], e[1]) in score_dict:
+                score_dict[e[0], e[1]] = 0.5 * (score_dict[e[0], e[1]] + si[j])
+            else:
+                score_dict[e[0], e[1]] = si[j]
+
+    nodes = np.sort(np.unique(edges.flatten()))
+    vertex = [Data(n) for n in nodes]
+    for (key, value) in score_dict.items():
+        vertex[key[0]].add_link(vertex[key[1]], value)
+
+    comps = connected_components(vertex, score_dict, th)
+
+    return comps
+
+
+def graph_propagation(edges, score, max_sz, step=0.1, beg_th=0.5, pool=None):
+
+    edges = np.sort(edges, axis=1)
+    th = score.min()
+    # th = beg_th
+    # construct graph
+    score_dict = {} # score lookup table
+    if pool is None:
+        for i,e in enumerate(edges):
+            score_dict[e[0], e[1]] = score[i]
+    elif pool == 'avg':
+        for i,e in enumerate(edges):
+            if (e[0], e[1]) in score_dict:
+                score_dict[e[0], e[1]] = 0.5*(score_dict[e[0], e[1]] + score[i])
+            else:
+                score_dict[e[0], e[1]] = score[i]
+
+    elif pool == 'max':
+        for i,e in enumerate(edges):
+            if (e[0],e[1]) in score_dict:
+                score_dict[e[0], e[1]] = max(score_dict[e[0], e[1]] , score[i])
+            else:
+                score_dict[e[0], e[1]] = score[i]
+    else:
+        raise ValueError('Pooling operation not supported')
+
+    nodes = np.sort(np.unique(edges.flatten()))
+    mapping = -1 * np.ones((nodes.max()+1), dtype=np.int)
+    mapping[nodes] = np.arange(nodes.shape[0])
+    link_idx = mapping[edges]
+    vertex = [Data(n) for n in nodes]
+    for l, s in zip(link_idx, score):
+        vertex[l[0]].add_link(vertex[l[1]], s)
+
+    # first iteration
+    comps, remain = connected_components_constraint(vertex, max_sz)
+
+    # iteration
+    components = comps[:]
+    while remain:
+        th = th + (1 - th) * step
+        comps, remain = connected_components_constraint(remain, max_sz, score_dict, th)
+        components.extend(comps)
+    return components
+
+
+def graph_propagation_soft(edges, score, max_sz, step=0.1, **kwargs):
+
+    edges = np.sort(edges, axis=1)
+    th = score.min()
+
+    # construct graph
+    score_dict = {} # score lookup table
+    for i,e in enumerate(edges):
+        score_dict[e[0], e[1]] = score[i]
+
+    nodes = np.sort(np.unique(edges.flatten()))
+    mapping = -1 * np.ones((nodes.max()+1), dtype=np.int)
+    mapping[nodes] = np.arange(nodes.shape[0])
+    link_idx = mapping[edges]
+    vertex = [Data(n) for n in nodes]
+    for l, s in zip(link_idx, score):
+        vertex[l[0]].add_link(vertex[l[1]], s)
+
+    # first iteration
+    comps, remain = connected_components_constraint(vertex, max_sz)
+    first_vertex_idx = np.array([mapping[n.name] for c in comps for n in c])
+    fusion_vertex_idx = np.setdiff1d(np.arange(nodes.shape[0]), first_vertex_idx, assume_unique=True)
+    # iteration
+    components = comps[:]
+    while remain:
+        th = th + (1 - th) * step
+        comps, remain = connected_components_constraint(remain, max_sz, score_dict, th)
+        components.extend(comps)
+    label_dict = {}
+    for i,c in enumerate(components):
+        for n in c:
+            label_dict[n.name] = i
+    print('Propagation ...')
+    prop_vertex = [vertex[idx] for idx in fusion_vertex_idx]
+    label, label_fusion = diffusion(prop_vertex, label_dict, score_dict, **kwargs)
+    return label, label_fusion
+
+
+def diffusion(vertex, label, score_dict, max_depth=5, weight_decay=0.6, normalize=True):
+    class BFSNode():
+        def __init__(self, node, depth, value):
+            self.node = node
+            self.depth = depth
+            self.value = value
+            
+    label_fusion = {}
+    for name in label.keys():
+        label_fusion[name] = {label[name]: 1.0}
+    prog = 0
+    prog_step = len(vertex) // 20
+    start = time.time()
+    for root in vertex:
+        if prog % prog_step == 0:
+            print("progress: {} / {}, elapsed time: {}".format(prog, len(vertex), time.time() - start))
+        prog += 1
+        #queue = {[root, 0, 1.0]}
+        queue = {BFSNode(root, 0, 1.0)}
+        visited = [root.name]
+        root_label = label[root.name]
+        while queue:
+            curr = queue.pop()
+            if curr.depth >= max_depth: # pruning
+                continue
+            neighbors = curr.node.links
+            tmp_value = []
+            tmp_neighbor = []
+            for n in neighbors:
+                if n.name not in visited:
+                    sub_value = score_dict[tuple(sorted([curr.node.name, n.name]))] * weight_decay * curr.value
+                    tmp_value.append(sub_value)
+                    tmp_neighbor.append(n)
+                    if root_label not in label_fusion[n.name].keys():
+                        label_fusion[n.name][root_label] = sub_value
+                    else:
+                        label_fusion[n.name][root_label] += sub_value
+                    visited.append(n.name)
+                    #queue.add([n, curr.depth+1, sub_value])
+            sortidx = np.argsort(tmp_value)[::-1]
+            for si in sortidx:
+                queue.add(BFSNode(tmp_neighbor[si], curr.depth+1, tmp_value[si]))
+    if normalize:
+        for name in label_fusion.keys():
+            summ = sum(label_fusion[name].values())
+            for k in label_fusion[name].keys():
+                label_fusion[name][k] /= summ
+    return label, label_fusion
+
+
+def clusters2labels(clusters, n_nodes):
+    labels = (-1)* np.ones((n_nodes,))
+    for ci, c in enumerate(clusters):
+        for xid in c:
+            labels[xid.name] = ci
+    assert np.sum(labels < 0) < 1
+    return labels
+
+
+def single_remove(bbox, pred):
+    single_idcs = np.zeros_like(pred)
+    pred_unique = np.unique(pred)
+    for u in pred_unique:
+        idcs = pred == u
+        if np.sum(idcs) == 1:
+            single_idcs[np.where(idcs)[0][0]] = 1
+    remain_idcs = [i for i in range(len(pred)) if not single_idcs[i]]
+    remain_idcs = np.asarray(remain_idcs)
+    return bbox[remain_idcs, :], pred[remain_idcs]
+

IndicPhotoOCR/detection/textbpn/util/io.py

@@ -0,0 +1,233 @@
+#coding=utf-8
+'''
+Created on 2016年9月27日
+
+@author: dengdan
+
+Tool  functions for file system operation and I/O. 
+In the style of linux shell commands
+'''
+import os
+import pickle as pkl
+import subprocess
+import logging
+from . import strs, io
+
+
+def mkdir(path):
+    """
+    If the target directory does not exists, it and its parent directories will created. 
+    """
+    path = get_absolute_path(path)
+    if not exists(path):
+        os.makedirs(path)
+    return path
+
+def make_parent_dir(path):
+    """make the parent directories for a file."""
+    parent_dir = get_dir(path)
+    mkdir(parent_dir)
+    
+    
+def pwd():
+    return os.getcwd()
+
+def dump(path, obj):
+    path = get_absolute_path(path)
+    parent_path = get_dir(path)
+    mkdir(parent_path)
+    with open(path, 'w') as f:
+        logging.info('dumping file:' + path);
+        pkl.dump(obj, f)
+
+def load(path):
+    path = get_absolute_path(path)
+    with open(path, 'r') as f:
+        data = pkl.load(f)
+    return data
+
+def join_path(a, *p):
+    return os.path.join(a, *p)
+
+def is_dir(path):
+    path = get_absolute_path(path)
+    return os.path.isdir(path)
+
+is_directory = is_dir
+
+def is_path(path):
+    path = get_absolute_path(path)
+    return os.path.ispath(path)
+    
+def get_dir(path):
+    '''
+    return the directory it belongs to.
+    if path is a directory itself, itself will be return 
+    '''
+    path = get_absolute_path(path)
+    if is_dir(path):
+        return path;
+    return os.path.split(path)[0]
+
+def get_parent_dir(path):
+    current_dir = get_dir(path)
+    return get_absolute_path(join_path(current_dir, '..'))
+
+def get_filename(path):
+    return os.path.split(path)[1]
+
+def get_absolute_path(p):
+    if p.startswith('~'):
+        p = os.path.expanduser(p)
+    return os.path.abspath(p)
+
+def cd(p):
+    p = get_absolute_path(p)
+    os.chdir(p)
+    
+def ls(path = '.', suffix = None):
+    """
+    list files in a directory.
+    return file names in a list
+    """
+    path = get_absolute_path(path)
+    files = os.listdir(path)
+
+    if suffix is None:       
+        return files
+        
+    filtered = []
+    for f in files:
+        if string.ends_with(f, suffix, ignore_case = True):
+            filtered.append(f)
+    
+    return filtered
+
+def find_files(pattern):
+    import glob
+    return glob.glob(pattern)
+
+def read_lines(p):
+    """return the text in a file in lines as a list """
+    p = get_absolute_path(p)
+    f = open(p,'r')
+    return f.readlines()
+    
+def write_lines(p, lines, append_break = False):
+    p = get_absolute_path(p)
+    make_parent_dir(p)
+    with open(p, 'w') as f:
+        for line in lines:
+            if append_break:
+                f.write(line + '\n')
+            else:
+                f.write(line)
+
+def cat(p):
+    """return the text in a file as a whole"""
+    cmd = 'cat ' + p
+    return subprocess.getoutput(cmd)
+
+def exists(path):
+    path = get_absolute_path(path)
+    return os.path.exists(path)
+
+def not_exists(path):
+    return not exists(path)
+
+def load_mat(path):
+    import scipy.io as sio # type: ignore
+    path = get_absolute_path(path)
+    return sio.loadmat(path)
+
+def dump_mat(path, dict_obj, append = True):
+    import scipy.io as sio # type: ignore
+    path = get_absolute_path(path)
+    make_parent_dir(path)
+    sio.savemat(file_name = path, mdict =  dict_obj, appendmat = append)
+    
+def dir_mat(path):
+    '''
+    list the variables in mat file.
+    return a list: [(name, shape, dtype), ...]
+    '''
+    import scipy.io as sio # type: ignore
+    path = get_absolute_path(path)
+    return sio.whosmat(path)
+    
+SIZE_UNIT_K = 1024
+SIZE_UNIT_M = SIZE_UNIT_K ** 2
+SIZE_UNIT_G = SIZE_UNIT_K ** 3
+def get_file_size(path, unit = SIZE_UNIT_K):
+    size = os.path.getsize(get_absolute_path(path))
+    return size * 1.0 / unit
+    
+    
+def create_h5(path):
+    import h5py # type: ignore
+    path = get_absolute_path(path)
+    make_parent_dir(path)
+    return h5py.File(path, 'w');
+
+def open_h5(path, mode = 'r'):
+    import h5py
+    path = get_absolute_path(path)
+    return h5py.File(path, mode);
+    
+def read_h5(h5, key):
+    return h5[key][:]
+def read_h5_attrs(h5, key, attrs):
+    return h5[key].attrs[attrs]
+    
+def copy(src, dest):
+    io.make_parent_dir(dest)
+    import shutil
+    shutil.copy(get_absolute_path(src), get_absolute_path(dest))
+    
+cp = copy
+
+def remove(p):
+    import os
+    os.remove(get_absolute_path(p))
+rm = remove
+
+def search(pattern, path, file_only = True):
+    """
+    Search files whose name matches the give pattern. The search scope
+    is the directory and sub-directories of 'path'. 
+    """
+    path = get_absolute_path(path)
+    pattern_here = io.join_path(path, pattern)
+    targets = []
+    
+    # find matchings in current directory
+    candidates = find_files(pattern_here)
+    for can in candidates:
+        if io.is_dir(can) and file_only:
+            continue
+        else:
+            targets.append(can)
+            
+    # find matching in sub-dirs
+    files = ls(path)
+    for f in files:
+        fpath = io.join_path(path, f)
+        if is_dir(fpath):
+            targets_in_sub_dir = search(pattern, fpath, file_only)
+            targets.extend(targets_in_sub_dir)
+    return targets
+
+def dump_json(path, data):
+    import ujson as json
+    path = get_absolute_path(path)
+    make_parent_dir(path)
+
+    with open(path, 'w') as f:
+        json.dump(data, f)
+    return path
+
+def load_json(path):
+    import ujson as json
+    path = get_absolute_path(path)
+    with open(path, 'r')  as f:
+        return json.load(f)