Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,176 @@
+# Output directories
+outputs/
+multirun/
+ray_results/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+# requirements/core.*.txt
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+.python-version
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# IDE
+.idea/
+
+########## CUSTOM FOLDER ##############
+README_original.md
+
+results/
+images
+bharatSTR/East/tmp
+bharatSTR/models
+bharatSTR/images
+__pycache__/
+bharatSTR/
+
+IndicPhotoOCR/detection/East
+IndicPhotoOCR/recognition/models
+
+IndicPhotoOCR/script_identification/images
+IndicPhotoOCR/script_identification/models
+
+
+build/
+dist/
+test.png
+static/pics/IndicPhotoOCR.gif
+input_image.jpg
+output_image.png
+flagged
+
+

CHANGELOG.md

@@ -0,0 +1,19 @@
+# Changelog
+
+## [1.1.0] - 2024-11-24
+### Added
+- Updated package naming convention
+
+## [1.0.3] - 2024-11-06
+### Added
+- Python package requirements sorted with setup.py
+
+## [1.0.2] - 2024-11-01
+### Added
+- Added language support for 10 additional models in the recognition module.
+
+## [1.0.1] - 2024-10-28
+### Added
+- Added support for detecting polygonal bounding boxes in `visualize_detection`.
+- Introduced the `show` argument in `visualize_detection` to control image display.
+

Dockerfile

@@ -0,0 +1,45 @@
+# Use NVIDIA PyTorch as the base image
+FROM nvcr.io/nvidia/pytorch:23.12-py3
+
+# Install additional dependencies
+RUN apt-get update && apt-get install -y ffmpeg libsm6 libxext6
+
+# Set environment variables for Miniconda and Conda environment
+ENV CONDA_DIR /opt/conda
+ENV PATH $CONDA_DIR/bin:$PATH
+
+# Install Miniconda
+RUN apt-get update && apt-get install -y wget && \
+    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p $CONDA_DIR && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+# Create a new Conda environment named "bocr" with Python 3.9
+RUN conda create -n bocr python=3.9 -y
+
+# Initialize conda 
+RUN conda init
+
+# Reload the env configs
+RUN source ~/.bashrc
+
+# Make RUN commands use the bocr environment
+SHELL ["conda", "run", "-n", "bocr", "/bin/bash", "-c"]
+
+# # Set default shell to bash
+# SHELL ["/bin/bash", "-c"]
+
+# # Clone BharatOCR repository
+# RUN git clone https://github.com/Bhashini-IITJ/BharatOCR.git && \
+#     git switch photoOCR && \
+#     cd IndicPhotoOCR && \
+#     python setup.py sdist bdist_wheel && \
+#     pip install ./dist/IndicPhotoOCR-1.1.0-py3-none-any.whl[cu118] --extra-index-url https://download.pytorch.org/whl/cu118
+
+# # # Set default command to run BharatOCR
+# CMD ["conda", "run", "-n", "bocr", "python", "-m", "IndicPhotoOCR.ocr"]
+
+
+# cd IndicPhotoOCR
+# sudo docker build -t indicphotoocr:latest .
+# sudo docker run --gpus all --rm -it indicphotoocr:latest

IndicPhotoOCR/detection/east_config.py

@@ -0,0 +1,39 @@
+
+# data-config
+import numpy as np
+
+train_data_path = './dataset/train/'
+train_batch_size_per_gpu = 14  # 14
+num_workers = 24  # 24
+gpu_ids = [0]  # [0,1,2,3]
+gpu = 1  # 4
+input_size = 512  # 预处理后归一化后图像尺寸
+background_ratio = 3. / 8  # 纯背景样本比例
+random_scale = np.array([0.5, 1, 2.0, 3.0])  # 提取多尺度图片信息
+geometry = 'RBOX'  # 选择使用几何特征图类型
+max_image_large_side = 1280
+max_text_size = 800
+min_text_size = 10
+min_crop_side_ratio = 0.1
+means=[100, 100, 100]
+pretrained = True  # 是否加载基础网络的预训练模型
+pretrained_basemodel_path = 'IndicPhotoOCR/detection/East/tmp/backbone_net/mobilenet_v2.pth.tar'
+pre_lr = 1e-4  # 基础网络的初始学习率
+lr = 1e-3  # 后面网络的初始学习率
+decay_steps = 50  # decayed_learning_rate = learning_rate * decay_rate ^ (global_epoch / decay_steps)
+decay_rate = 0.97
+init_type = 'xavier'  # 网络参数初始化方式
+resume = True  # 整体网络是否恢复原来保存的模型
+checkpoint = 'IndicPhotoOCR/detection/East/tmp/epoch_990_checkpoint.pth.tar'  # 指定具体路径及文件名
+max_epochs = 1000  # 最大迭代epochs数
+l2_weight_decay = 1e-6  # l2正则化惩罚项权重
+print_freq = 10  # 每10个batch输出损失结果
+save_eval_iteration = 50  # 每10个epoch保存一次模型,并做一次评价
+save_model_path = './tmp/'  # 模型保存路径
+test_img_path = './dataset/full_set'  # demo测试样本路径'./demo/test_img/'，数据集测试为'./dataset/test/'
+res_img_path = 'results'  # demo结果存放路径'./demo/result_img/'，数据集测试为 './dataset/test_result/'
+write_images = True  # 是否输出图像结果
+score_map_thresh = 0.8  # 置信度阈值
+box_thresh = 0.1  # 文本框中置信度平均值的阈值
+nms_thres = 0.2  # 局部非极大抑制IOU阈值
+compute_hmean_path = './dataset/test_compute_hmean/'

IndicPhotoOCR/detection/east_detector.py

@@ -0,0 +1,87 @@
+import os
+import torch
+import cv2
+import numpy as np
+import time
+import warnings
+
+
+import IndicPhotoOCR.detection.east_config as cfg
+from IndicPhotoOCR.detection.east_utils import ModelManager
+from IndicPhotoOCR.detection.east_model import East
+import IndicPhotoOCR.detection.east_utils as utils
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+class EASTdetector:
+    def __init__(self, model_name= "east", model_path=None):
+        self.model_path = model_path
+        # self.model_manager = ModelManager()
+        # self.model_manager.ensure_model(model_name)
+        # self.ensure_model(self.model_name)
+        # self.root_model_dir = "BharatSTR/bharatOCR/detection/East/tmp"
+
+    def detect(self, image_path, model_checkpoint, device):
+        # Load image
+        im = cv2.imread(image_path)
+        # im = cv2.imread(image_path)[:, :, ::-1]
+
+        # Initialize the EAST model and load checkpoint
+        model = East()
+        model = torch.nn.DataParallel(model, device_ids=cfg.gpu_ids)
+
+        # Load the model checkpoint with weights_only=True
+        checkpoint = torch.load(model_checkpoint, map_location=torch.device(device), weights_only=True)
+        model.load_state_dict(checkpoint['state_dict'])
+        model.eval()
+
+        # Resize image and convert to tensor format
+        im_resized, (ratio_h, ratio_w) = utils.resize_image(im)
+        im_resized = im_resized.astype(np.float32).transpose(2, 0, 1)
+        im_tensor = torch.from_numpy(im_resized).unsqueeze(0).cpu()
+
+        # Inference
+        timer = {'net': 0, 'restore': 0, 'nms': 0}
+        start = time.time()
+        score, geometry = model(im_tensor)
+        timer['net'] = time.time() - start
+
+        # Process output
+        score = score.permute(0, 2, 3, 1).data.cpu().numpy()
+        geometry = geometry.permute(0, 2, 3, 1).data.cpu().numpy()
+        
+        # Detect boxes
+        boxes, timer = utils.detect(
+            score_map=score, geo_map=geometry, timer=timer,
+            score_map_thresh=cfg.score_map_thresh, box_thresh=cfg.box_thresh,
+            nms_thres=cfg.box_thresh
+        )
+        bbox_result_dict = {'detections': []}
+
+        # Parse detected boxes and adjust coordinates
+        if boxes is not None:
+            boxes = boxes[:, :8].reshape((-1, 4, 2))
+            boxes[:, :, 0] /= ratio_w
+            boxes[:, :, 1] /= ratio_h
+            for box in boxes:
+                box = utils.sort_poly(box.astype(np.int32))
+                if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm(box[3] - box[0]) < 5:
+                    continue
+                bbox_result_dict['detections'].append([
+                    [int(coord[0]), int(coord[1])] for coord in box
+                ])
+
+        return bbox_result_dict
+
+# 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_checkpoint', type=str, required=True, help='Path to the model checkpoint file')
+#     args = parser.parse_args()
+
+#     # Run prediction and get results as dictionary
+#     detection_result = predict(args.image_path, args.device, args.model_checkpoint)
+#     print(detection_result)

IndicPhotoOCR/detection/east_locality_aware_nms.py

@@ -0,0 +1,75 @@
+
+import numpy as np
+from shapely.geometry import Polygon
+
+
+def intersection(g, p):
+    g = Polygon(g[:8].reshape((4, 2)))
+    p = Polygon(p[:8].reshape((4, 2)))
+    if not g.is_valid or not p.is_valid:
+        return 0
+    inter = Polygon(g).intersection(Polygon(p)).area
+    union = g.area + p.area - inter
+    if union == 0:
+        return 0
+    else:
+        return inter/union
+
+
+def weighted_merge(g, p):
+    # g[0]=min(g[0],p[0])
+    # g[1] = min(g[1], p[1])
+    # g[4] = max(g[4], p[4])
+    # g[5]= max(g[5],p[5])
+    #
+    # g[2] = max(g[2], p[2])
+    # g[3] = min(g[3], p[3])
+    # g[6] = min(g[6], p[6])
+    # g[7] = max(g[7], p[7])
+
+    g[:8] = (g[8] * g[:8] + p[8] * p[:8])/(g[8] + p[8])
+    g[8] = (g[8] + p[8])
+    return g
+
+
+def standard_nms(S, thres):
+    order = np.argsort(S[:, 8])[::-1]
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        ovr = np.array([intersection(S[i], S[t]) for t in order[1:]])
+
+        inds = np.where(ovr <= thres)[0]
+        order = order[inds+1]
+
+    return S[keep]
+
+
+def nms_locality(polys, thres=0.3):
+    '''
+    locality aware nms of EAST
+    :param polys: a N*9 numpy array. first 8 coordinates, then prob
+    :return: boxes after nms
+    '''
+    S = []
+    p = None
+    for g in polys:
+        if p is not None and intersection(g, p) > thres:
+            p = weighted_merge(g, p)
+        else:
+            if p is not None:
+                S.append(p)
+            p = g
+    if p is not None:
+        S.append(p)
+
+    if len(S) == 0:
+        return np.array([])
+    return standard_nms(np.array(S), thres)
+
+
+if __name__ == '__main__':
+    # 343,350,448,135,474,143,369,359
+    print(Polygon(np.array([[343, 350], [448, 135],
+                            [474, 143], [369, 359]])).area)

IndicPhotoOCR/detection/east_model.py

@@ -0,0 +1,242 @@
+
+import torch.nn as nn
+import math
+import torch
+
+
+from IndicPhotoOCR.detection import east_config as cfg
+from IndicPhotoOCR.detection import east_utils as utils
+
+
+def conv_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, width_mult=1.):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            [6, 160, 3, 2],
+            # [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        # assert input_size % 32 == 0
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
+        self.features = [conv_bn(3, input_channel, 2)]
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
+                input_channel = output_channel
+
+        # make it nn.Sequential
+        self.features = nn.Sequential(*self.features)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        # x = x.mean(3).mean(2)
+        # x = self.classifier(x)
+        return x
+
+    def _initialize_weights(self):
+        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))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                n = m.weight.size(1)
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+
+def mobilenet(pretrained=True, **kwargs):
+    """
+    Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = MobileNetV2()
+    if pretrained:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(cfg.pretrained_basemodel_path,map_location=torch.device('cpu'), weights_only=True)
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        # state_dict = torch.load(cfg.pretrained_basemodel_path)  # add map_location='cpu' if no gpu
+        # model.load_state_dict(state_dict)
+
+    return model
+
+
+class East(nn.Module):
+    def __init__(self):
+        super(East, self).__init__()
+        self.mobilenet = mobilenet(True)
+        # self.si for stage i
+        self.s1 = nn.Sequential(*list(self.mobilenet.children())[0][0:4])
+        self.s2 = nn.Sequential(*list(self.mobilenet.children())[0][4:7])
+        self.s3 = nn.Sequential(*list(self.mobilenet.children())[0][7:14])
+        self.s4 = nn.Sequential(*list(self.mobilenet.children())[0][14:17])
+
+        self.conv1 = nn.Conv2d(160+96, 128, 1)
+        self.bn1 = nn.BatchNorm2d(128)
+        self.relu1 = nn.ReLU()
+
+        self.conv2 = nn.Conv2d(128, 128, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(128)
+        self.relu2 = nn.ReLU()
+
+        self.conv3 = nn.Conv2d(128+32, 64, 1)
+        self.bn3 = nn.BatchNorm2d(64)
+        self.relu3 = nn.ReLU()
+
+        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
+        self.bn4 = nn.BatchNorm2d(64)
+        self.relu4 = nn.ReLU()
+
+        self.conv5 = nn.Conv2d(64+24, 64, 1)
+        self.bn5 = nn.BatchNorm2d(64)
+        self.relu5 = nn.ReLU()
+
+        self.conv6 = nn.Conv2d(64, 32, 3, padding=1)
+        self.bn6 = nn.BatchNorm2d(32)
+        self.relu6 = nn.ReLU()
+
+        self.conv7 = nn.Conv2d(32, 32, 3, padding=1)
+        self.bn7 = nn.BatchNorm2d(32)
+        self.relu7 = nn.ReLU()
+
+        self.conv8 = nn.Conv2d(32, 1, 1)
+        self.sigmoid1 = nn.Sigmoid()
+        self.conv9 = nn.Conv2d(32, 4, 1)
+        self.sigmoid2 = nn.Sigmoid()
+        self.conv10 = nn.Conv2d(32, 1, 1)
+        self.sigmoid3 = nn.Sigmoid()
+        self.unpool1 = nn.Upsample(scale_factor=2, mode='bilinear')
+        self.unpool2 = nn.Upsample(scale_factor=2, mode='bilinear')
+        self.unpool3 = nn.Upsample(scale_factor=2, mode='bilinear')
+
+        # utils.init_weights([self.conv1,self.conv2,self.conv3,self.conv4,
+        #                        self.conv5,self.conv6,self.conv7,self.conv8,
+        #                        self.conv9,self.conv10,self.bn1,self.bn2,
+        #                        self.bn3,self.bn4,self.bn5,self.bn6,self.bn7])
+
+    def forward(self, images):
+        images = utils.mean_image_subtraction(images)
+
+        f0 = self.s1(images)
+        f1 = self.s2(f0)
+        f2 = self.s3(f1)
+        f3 = self.s4(f2)
+
+        # _, f = self.mobilenet(images)
+        h = f3  # bs 2048 w/32 h/32
+        g = (self.unpool1(h))  # bs 2048 w/16 h/16
+        c = self.conv1(torch.cat((g, f2), 1))
+        c = self.bn1(c)
+        c = self.relu1(c)
+
+        h = self.conv2(c)  # bs 128 w/16 h/16
+        h = self.bn2(h)
+        h = self.relu2(h)
+        g = self.unpool2(h)  # bs 128 w/8 h/8
+        c = self.conv3(torch.cat((g, f1), 1))
+        c = self.bn3(c)
+        c = self.relu3(c)
+
+        h = self.conv4(c)  # bs 64 w/8 h/8
+        h = self.bn4(h)
+        h = self.relu4(h)
+        g = self.unpool3(h)  # bs 64 w/4 h/4
+        c = self.conv5(torch.cat((g, f0), 1))
+        c = self.bn5(c)
+        c = self.relu5(c)
+
+        h = self.conv6(c)  # bs 32 w/4 h/4
+        h = self.bn6(h)
+        h = self.relu6(h)
+        g = self.conv7(h)  # bs 32 w/4 h/4
+        g = self.bn7(g)
+        g = self.relu7(g)
+
+        F_score = self.conv8(g)  # bs 1 w/4 h/4
+        F_score = self.sigmoid1(F_score)
+        geo_map = self.conv9(g)
+        geo_map = self.sigmoid2(geo_map) * 512
+        angle_map = self.conv10(g)
+        angle_map = self.sigmoid3(angle_map)
+        angle_map = (angle_map - 0.5) * math.pi / 2
+
+        F_geometry = torch.cat((geo_map, angle_map), 1)  # bs 5 w/4 h/4
+
+        return F_score, F_geometry
+
+
+model=East()

IndicPhotoOCR/detection/east_preprossing.py

@@ -0,0 +1,681 @@
+
+# coding:utf-8
+import glob
+import csv
+import cv2
+import os
+import numpy as np
+from shapely.geometry import Polygon
+
+
+from IndicPhotoOCR.detection import east_config as cfg
+from IndicPhotoOCR.detection import east_utils
+
+
+def get_images(img_root):
+    files = []
+    for ext in ['jpg']:
+        files.extend(glob.glob(
+            os.path.join(img_root, '*.{}'.format(ext))))
+        # print(glob.glob(
+        #     os.path.join(FLAGS.training_data_path, '*.{}'.format(ext))))
+    return files
+
+
+def load_annoataion(p):
+    '''
+    load annotation from the text file
+    :param p:
+    :return:
+    '''
+    text_polys = []
+    text_tags = []
+    if not os.path.exists(p):
+        return np.array(text_polys, dtype=np.float32)
+    with open(p, 'r', encoding='UTF-8') as f:
+        reader = csv.reader(f)
+        for line in reader:
+            label = line[-1]
+            # strip BOM. \ufeff for python3,  \xef\xbb\bf for python2
+            line = [i.strip('\ufeff').strip('\xef\xbb\xbf') for i in line]
+
+            x1, y1, x2, y2, x3, y3, x4, y4 = list(map(float, line[:8]))
+            text_polys.append([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])
+            # print(text_polys)
+            if label == '*' or label == '###':
+                text_tags.append(True)
+            else:
+                text_tags.append(False)
+        return np.array(text_polys, dtype=np.float32), np.array(text_tags, dtype=np.bool)
+
+
+def polygon_area(poly):
+    '''
+    compute area of a polygon
+    :param poly:
+    :return:
+    '''
+    edge = [
+        (poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
+        (poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
+        (poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
+        (poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1])
+    ]
+    return np.sum(edge) / 2.
+
+
+def check_and_validate_polys(polys, tags, xxx_todo_changeme):
+    '''
+    check so that the text poly is in the same direction,
+    and also filter some invalid polygons
+    :param polys:
+    :param tags:
+    :return:
+    '''
+    (h, w) = xxx_todo_changeme
+    if polys.shape[0] == 0:
+        return polys
+    polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
+    polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
+
+    validated_polys = []
+    validated_tags = []
+
+    # 判断四边形的点时针方向，以及是否是有效四边形
+    for poly, tag in zip(polys, tags):
+        p_area = polygon_area(poly)
+        if abs(p_area) < 1:
+            # print poly
+            print('invalid poly')
+            continue
+        if p_area > 0:
+            print('poly in wrong direction')
+            poly = poly[(0, 3, 2, 1), :]
+        validated_polys.append(poly)
+        validated_tags.append(tag)
+    return np.array(validated_polys), np.array(validated_tags)
+
+
+def crop_area(im, polys, tags, crop_background=False, max_tries=100):
+    '''
+    make random crop from the input image
+    :param im:
+    :param polys:
+    :param tags:
+    :param crop_background:
+    :param max_tries:
+    :return:
+    '''
+    h, w, _ = im.shape
+    pad_h = h // 10
+    pad_w = w // 10
+    h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
+    w_array = np.zeros((w + pad_w * 2), 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]
+    if len(h_axis) == 0 or len(w_axis) == 0:
+        return im, polys, tags
+    for i in range(max_tries):  # 试验50次
+        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 < cfg.min_crop_side_ratio * w or ymax - ymin < cfg.min_crop_side_ratio * h:
+            # area too small
+            continue
+        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 = []
+        if len(selected_polys) == 0:
+            # no text in this area
+            if crop_background:
+                return im[ymin:ymax + 1, xmin:xmax + 1, :], polys[selected_polys], tags[selected_polys]
+            else:
+                continue
+        im = im[ymin:ymax + 1, xmin:xmax + 1, :]
+        polys = polys[selected_polys]
+        tags = tags[selected_polys]
+        polys[:, :, 0] -= xmin
+        polys[:, :, 1] -= ymin
+        return im, polys, tags
+
+    return im, polys, tags
+
+
+def shrink_poly(poly, r):
+    '''
+    fit a poly inside the origin poly, maybe bugs here...
+    used for generate the score map
+    :param poly: the text poly
+    :param r: r in the paper
+    :return: the shrinked poly
+    '''
+    # shrink ratio
+    R = 0.3
+    # find the longer pair
+    if np.linalg.norm(poly[0] - poly[1]) + np.linalg.norm(poly[2] - poly[3]) > \
+                    np.linalg.norm(poly[0] - poly[3]) + np.linalg.norm(poly[1] - poly[2]):
+        # first move (p0, p1), (p2, p3), then (p0, p3), (p1, p2)
+        ## p0, p1
+        theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
+        poly[0][0] += R * r[0] * np.cos(theta)
+        poly[0][1] += R * r[0] * np.sin(theta)
+        poly[1][0] -= R * r[1] * np.cos(theta)
+        poly[1][1] -= R * r[1] * np.sin(theta)
+        ## p2, p3
+        theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
+        poly[3][0] += R * r[3] * np.cos(theta)
+        poly[3][1] += R * r[3] * np.sin(theta)
+        poly[2][0] -= R * r[2] * np.cos(theta)
+        poly[2][1] -= R * r[2] * np.sin(theta)
+        ## p0, p3
+        theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
+        poly[0][0] += R * r[0] * np.sin(theta)
+        poly[0][1] += R * r[0] * np.cos(theta)
+        poly[3][0] -= R * r[3] * np.sin(theta)
+        poly[3][1] -= R * r[3] * np.cos(theta)
+        ## p1, p2
+        theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
+        poly[1][0] += R * r[1] * np.sin(theta)
+        poly[1][1] += R * r[1] * np.cos(theta)
+        poly[2][0] -= R * r[2] * np.sin(theta)
+        poly[2][1] -= R * r[2] * np.cos(theta)
+    else:
+        ## p0, p3
+        # print poly
+        theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
+        poly[0][0] += R * r[0] * np.sin(theta)
+        poly[0][1] += R * r[0] * np.cos(theta)
+        poly[3][0] -= R * r[3] * np.sin(theta)
+        poly[3][1] -= R * r[3] * np.cos(theta)
+        ## p1, p2
+        theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
+        poly[1][0] += R * r[1] * np.sin(theta)
+        poly[1][1] += R * r[1] * np.cos(theta)
+        poly[2][0] -= R * r[2] * np.sin(theta)
+        poly[2][1] -= R * r[2] * np.cos(theta)
+        ## p0, p1
+        theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
+        poly[0][0] += R * r[0] * np.cos(theta)
+        poly[0][1] += R * r[0] * np.sin(theta)
+        poly[1][0] -= R * r[1] * np.cos(theta)
+        poly[1][1] -= R * r[1] * np.sin(theta)
+        ## p2, p3
+        theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
+        poly[3][0] += R * r[3] * np.cos(theta)
+        poly[3][1] += R * r[3] * np.sin(theta)
+        poly[2][0] -= R * r[2] * np.cos(theta)
+        poly[2][1] -= R * r[2] * np.sin(theta)
+    return poly
+
+
+# def point_dist_to_line(p1, p2, p3):
+#     # compute the distance from p3 to p1-p2
+#     return np.linalg.norm(np.cross(p2 - p1, p1 - p3)) / np.linalg.norm(p2 - p1)
+
+
+# 点p3到直线p12的距离
+def point_dist_to_line(p1, p2, p3):
+    # compute the distance from p3 to p1-p2
+    # return np.linalg.norm(np.cross(p2 - p1, p1 - p3)) / np.linalg.norm(p2 - p1)
+    a = np.linalg.norm(p1 - p2)
+    b = np.linalg.norm(p2 - p3)
+    c = np.linalg.norm(p3 - p1)
+    s = (a + b + c) / 2.0
+    area = np.abs((s * (s - a) * (s - b) * (s - c))) ** 0.5
+    if a < 1.0:
+        return (b + c) / 2.0
+    return 2 * area / a
+
+
+def fit_line(p1, p2):
+    # fit a line ax+by+c = 0
+    if p1[0] == p1[1]:
+        return [1., 0., -p1[0]]
+    else:
+        [k, b] = np.polyfit(p1, p2, deg=1)
+        return [k, -1., b]
+
+
+def line_cross_point(line1, line2):
+    # line1 0= ax+by+c, compute the cross point of line1 and line2
+    if line1[0] != 0 and line1[0] == line2[0]:
+        print('Cross point does not exist')
+        return None
+    if line1[0] == 0 and line2[0] == 0:
+        print('Cross point does not exist')
+        return None
+    if line1[1] == 0:
+        x = -line1[2]
+        y = line2[0] * x + line2[2]
+    elif line2[1] == 0:
+        x = -line2[2]
+        y = line1[0] * x + line1[2]
+    else:
+        k1, _, b1 = line1
+        k2, _, b2 = line2
+        x = -(b1 - b2) / (k1 - k2)
+        y = k1 * x + b1
+    return np.array([x, y], dtype=np.float32)
+
+
+def line_verticle(line, point):
+    # get the verticle line from line across point
+    if line[1] == 0:
+        verticle = [0, -1, point[1]]
+    else:
+        if line[0] == 0:
+            verticle = [1, 0, -point[0]]
+        else:
+            verticle = [-1. / line[0], -1, point[1] - (-1 / line[0] * point[0])]
+    return verticle
+
+
+def rectangle_from_parallelogram(poly):
+    '''
+    fit a rectangle from a parallelogram
+    :param poly:
+    :return:
+    '''
+    p0, p1, p2, p3 = poly
+    angle_p0 = np.arccos(np.dot(p1 - p0, p3 - p0) / (np.linalg.norm(p0 - p1) * np.linalg.norm(p3 - p0)))
+    if angle_p0 < 0.5 * np.pi:
+        if np.linalg.norm(p0 - p1) > np.linalg.norm(p0 - p3):
+            # p0 and p2
+            ## p0
+            p2p3 = fit_line([p2[0], p3[0]], [p2[1], p3[1]])
+            p2p3_verticle = line_verticle(p2p3, p0)
+
+            new_p3 = line_cross_point(p2p3, p2p3_verticle)
+            ## p2
+            p0p1 = fit_line([p0[0], p1[0]], [p0[1], p1[1]])
+            p0p1_verticle = line_verticle(p0p1, p2)
+
+            new_p1 = line_cross_point(p0p1, p0p1_verticle)
+            return np.array([p0, new_p1, p2, new_p3], dtype=np.float32)
+        else:
+            p1p2 = fit_line([p1[0], p2[0]], [p1[1], p2[1]])
+            p1p2_verticle = line_verticle(p1p2, p0)
+
+            new_p1 = line_cross_point(p1p2, p1p2_verticle)
+            p0p3 = fit_line([p0[0], p3[0]], [p0[1], p3[1]])
+            p0p3_verticle = line_verticle(p0p3, p2)
+
+            new_p3 = line_cross_point(p0p3, p0p3_verticle)
+            return np.array([p0, new_p1, p2, new_p3], dtype=np.float32)
+    else:
+        if np.linalg.norm(p0 - p1) > np.linalg.norm(p0 - p3):
+            # p1 and p3
+            ## p1
+            p2p3 = fit_line([p2[0], p3[0]], [p2[1], p3[1]])
+            p2p3_verticle = line_verticle(p2p3, p1)
+
+            new_p2 = line_cross_point(p2p3, p2p3_verticle)
+            ## p3
+            p0p1 = fit_line([p0[0], p1[0]], [p0[1], p1[1]])
+            p0p1_verticle = line_verticle(p0p1, p3)
+
+            new_p0 = line_cross_point(p0p1, p0p1_verticle)
+            return np.array([new_p0, p1, new_p2, p3], dtype=np.float32)
+        else:
+            p0p3 = fit_line([p0[0], p3[0]], [p0[1], p3[1]])
+            p0p3_verticle = line_verticle(p0p3, p1)
+
+            new_p0 = line_cross_point(p0p3, p0p3_verticle)
+            p1p2 = fit_line([p1[0], p2[0]], [p1[1], p2[1]])
+            p1p2_verticle = line_verticle(p1p2, p3)
+
+            new_p2 = line_cross_point(p1p2, p1p2_verticle)
+            return np.array([new_p0, p1, new_p2, p3], dtype=np.float32)
+
+
+def sort_rectangle(poly):
+    # sort the four coordinates of the polygon, points in poly should be sorted clockwise
+    # First find the lowest point
+    p_lowest = np.argmax(poly[:, 1])
+    if np.count_nonzero(poly[:, 1] == poly[p_lowest, 1]) == 2:
+        # 底边平行于X轴, 那么p0为左上角 - if the bottom line is parallel to x-axis, then p0 must be the upper-left corner
+        p0_index = np.argmin(np.sum(poly, axis=1))
+        p1_index = (p0_index + 1) % 4
+        p2_index = (p0_index + 2) % 4
+        p3_index = (p0_index + 3) % 4
+        return poly[[p0_index, p1_index, p2_index, p3_index]], 0.
+    else:
+        # 找到最低点右边的点 - find the point that sits right to the lowest point
+        p_lowest_right = (p_lowest - 1) % 4
+        p_lowest_left = (p_lowest + 1) % 4
+        angle = np.arctan(
+            -(poly[p_lowest][1] - poly[p_lowest_right][1]) / (poly[p_lowest][0] - poly[p_lowest_right][0]))
+        # assert angle > 0
+        if angle <= 0:
+            print(angle, poly[p_lowest], poly[p_lowest_right])
+        if angle / np.pi * 180 > 45:
+            # 这个点为p2 - this point is p2
+            p2_index = p_lowest
+            p1_index = (p2_index - 1) % 4
+            p0_index = (p2_index - 2) % 4
+            p3_index = (p2_index + 1) % 4
+            return poly[[p0_index, p1_index, p2_index, p3_index]], -(np.pi / 2 - angle)
+        else:
+            # 这个点为p3 - this point is p3
+            p3_index = p_lowest
+            p0_index = (p3_index + 1) % 4
+            p1_index = (p3_index + 2) % 4
+            p2_index = (p3_index + 3) % 4
+            return poly[[p0_index, p1_index, p2_index, p3_index]], angle
+
+
+def restore_rectangle_rbox(origin, geometry):
+    d = geometry[:, :4]
+    angle = geometry[:, 4]
+    # for angle > 0
+    origin_0 = origin[angle >= 0]
+    d_0 = d[angle >= 0]
+    angle_0 = angle[angle >= 0]
+    if origin_0.shape[0] > 0:
+        p = np.array([np.zeros(d_0.shape[0]), -d_0[:, 0] - d_0[:, 2],
+                      d_0[:, 1] + d_0[:, 3], -d_0[:, 0] - d_0[:, 2],
+                      d_0[:, 1] + d_0[:, 3], np.zeros(d_0.shape[0]),
+                      np.zeros(d_0.shape[0]), np.zeros(d_0.shape[0]),
+                      d_0[:, 3], -d_0[:, 2]])
+        p = p.transpose((1, 0)).reshape((-1, 5, 2))  # N*5*2
+
+        rotate_matrix_x = np.array([np.cos(angle_0), np.sin(angle_0)]).transpose((1, 0))
+        rotate_matrix_x = np.repeat(rotate_matrix_x, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))  # N*5*2
+
+        rotate_matrix_y = np.array([-np.sin(angle_0), np.cos(angle_0)]).transpose((1, 0))
+        rotate_matrix_y = np.repeat(rotate_matrix_y, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))
+
+        p_rotate_x = np.sum(rotate_matrix_x * p, axis=2)[:, :, np.newaxis]  # N*5*1
+        p_rotate_y = np.sum(rotate_matrix_y * p, axis=2)[:, :, np.newaxis]  # N*5*1
+
+        p_rotate = np.concatenate([p_rotate_x, p_rotate_y], axis=2)  # N*5*2
+
+        p3_in_origin = origin_0 - p_rotate[:, 4, :]
+        new_p0 = p_rotate[:, 0, :] + p3_in_origin  # N*2
+        new_p1 = p_rotate[:, 1, :] + p3_in_origin
+        new_p2 = p_rotate[:, 2, :] + p3_in_origin
+        new_p3 = p_rotate[:, 3, :] + p3_in_origin
+
+        new_p_0 = np.concatenate([new_p0[:, np.newaxis, :], new_p1[:, np.newaxis, :],
+                                  new_p2[:, np.newaxis, :], new_p3[:, np.newaxis, :]], axis=1)  # N*4*2
+    else:
+        new_p_0 = np.zeros((0, 4, 2))
+    # for angle < 0
+    origin_1 = origin[angle < 0]
+    d_1 = d[angle < 0]
+    angle_1 = angle[angle < 0]
+    if origin_1.shape[0] > 0:
+        p = np.array([-d_1[:, 1] - d_1[:, 3], -d_1[:, 0] - d_1[:, 2],
+                      np.zeros(d_1.shape[0]), -d_1[:, 0] - d_1[:, 2],
+                      np.zeros(d_1.shape[0]), np.zeros(d_1.shape[0]),
+                      -d_1[:, 1] - d_1[:, 3], np.zeros(d_1.shape[0]),
+                      -d_1[:, 1], -d_1[:, 2]])
+        p = p.transpose((1, 0)).reshape((-1, 5, 2))  # N*5*2
+
+        rotate_matrix_x = np.array([np.cos(-angle_1), -np.sin(-angle_1)]).transpose((1, 0))
+        rotate_matrix_x = np.repeat(rotate_matrix_x, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))  # N*5*2
+
+        rotate_matrix_y = np.array([np.sin(-angle_1), np.cos(-angle_1)]).transpose((1, 0))
+        rotate_matrix_y = np.repeat(rotate_matrix_y, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))
+
+        p_rotate_x = np.sum(rotate_matrix_x * p, axis=2)[:, :, np.newaxis]  # N*5*1
+        p_rotate_y = np.sum(rotate_matrix_y * p, axis=2)[:, :, np.newaxis]  # N*5*1
+
+        p_rotate = np.concatenate([p_rotate_x, p_rotate_y], axis=2)  # N*5*2
+
+        p3_in_origin = origin_1 - p_rotate[:, 4, :]
+        new_p0 = p_rotate[:, 0, :] + p3_in_origin  # N*2
+        new_p1 = p_rotate[:, 1, :] + p3_in_origin
+        new_p2 = p_rotate[:, 2, :] + p3_in_origin
+        new_p3 = p_rotate[:, 3, :] + p3_in_origin
+
+        new_p_1 = np.concatenate([new_p0[:, np.newaxis, :], new_p1[:, np.newaxis, :],
+                                  new_p2[:, np.newaxis, :], new_p3[:, np.newaxis, :]], axis=1)  # N*4*2
+    else:
+        new_p_1 = np.zeros((0, 4, 2))
+    return np.concatenate([new_p_0, new_p_1])
+
+
+def restore_rectangle(origin, geometry):
+    return restore_rectangle_rbox(origin, geometry)
+
+
+def generate_rbox(im_size, polys, tags):
+    h, w = im_size
+    poly_mask = np.zeros((h, w), dtype=np.uint8)
+    score_map = np.zeros((h, w), dtype=np.uint8)
+    geo_map = np.zeros((h, w, 5), dtype=np.float32)
+    # mask used during traning, to ignore some hard areas，用于忽略那些过小的文本
+    training_mask = np.ones((h, w), dtype=np.uint8)
+    for poly_idx, poly_tag in enumerate(zip(polys, tags)):
+        poly = poly_tag[0]
+        tag = poly_tag[1]
+
+        # 对每个顶点，找到经过他的两条边中较短的那条
+        r = [None, None, None, None]
+        for i in range(4):
+            r[i] = min(np.linalg.norm(poly[i] - poly[(i + 1) % 4]),
+                       np.linalg.norm(poly[i] - poly[(i - 1) % 4]))
+        # score map
+        # 放缩边框为之前的0.3倍，并对边框对应score图中的位置进行填充
+        shrinked_poly = shrink_poly(poly.copy(), r).astype(np.int32)[np.newaxis, :, :]
+        cv2.fillPoly(score_map, shrinked_poly, 1)
+        cv2.fillPoly(poly_mask, shrinked_poly, poly_idx + 1)
+        # if the poly is too small, then ignore it during training
+        # 如果文本框标签太小或者txt中没具体标记是什么内容，即*或者###，则加掩模，训练时忽略该部分
+        poly_h = min(np.linalg.norm(poly[0] - poly[3]), np.linalg.norm(poly[1] - poly[2]))
+        poly_w = min(np.linalg.norm(poly[0] - poly[1]), np.linalg.norm(poly[2] - poly[3]))
+        if min(poly_h, poly_w) < cfg.min_text_size:
+            cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
+        if tag:
+            cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
+
+        # 当前新加入的文本框区域像素点
+        xy_in_poly = np.argwhere(poly_mask == (poly_idx + 1))
+        # if geometry == 'RBOX':
+        # 对任意两个顶点的组合生成一个平行四边形 - generate a parallelogram for any combination of two vertices
+        fitted_parallelograms = []
+        for i in range(4):
+            # 选中p0和p1的连线边，生成两个平行四边形
+            p0 = poly[i]
+            p1 = poly[(i + 1) % 4]
+            p2 = poly[(i + 2) % 4]
+            p3 = poly[(i + 3) % 4]
+            # 拟合ax+by+c=0
+            edge = fit_line([p0[0], p1[0]], [p0[1], p1[1]])
+            backward_edge = fit_line([p0[0], p3[0]], [p0[1], p3[1]])
+            forward_edge = fit_line([p1[0], p2[0]], [p1[1], p2[1]])
+            # 通过另外两个点距离edge的距离，来决定edge对应的平行线应该过p2还是p3
+            if point_dist_to_line(p0, p1, p2) > point_dist_to_line(p0, p1, p3):
+                # 平行线经过p2 - parallel lines through p2
+                if edge[1] == 0:
+                    edge_opposite = [1, 0, -p2[0]]
+                else:
+                    edge_opposite = [edge[0], -1, p2[1] - edge[0] * p2[0]]
+            else:
+                # 经过p3 - after p3
+                if edge[1] == 0:
+                    edge_opposite = [1, 0, -p3[0]]
+                else:
+                    edge_opposite = [edge[0], -1, p3[1] - edge[0] * p3[0]]
+            # move forward edge
+            new_p0 = p0
+            new_p1 = p1
+            new_p2 = p2
+            new_p3 = p3
+            new_p2 = line_cross_point(forward_edge, edge_opposite)
+            if point_dist_to_line(p1, new_p2, p0) > point_dist_to_line(p1, new_p2, p3):
+                # across p0
+                if forward_edge[1] == 0:
+                    forward_opposite = [1, 0, -p0[0]]
+                else:
+                    forward_opposite = [forward_edge[0], -1, p0[1] - forward_edge[0] * p0[0]]
+            else:
+                # across p3
+                if forward_edge[1] == 0:
+                    forward_opposite = [1, 0, -p3[0]]
+                else:
+                    forward_opposite = [forward_edge[0], -1, p3[1] - forward_edge[0] * p3[0]]
+            new_p0 = line_cross_point(forward_opposite, edge)
+            new_p3 = line_cross_point(forward_opposite, edge_opposite)
+            fitted_parallelograms.append([new_p0, new_p1, new_p2, new_p3, new_p0])
+            # or move backward edge
+            new_p0 = p0
+            new_p1 = p1
+            new_p2 = p2
+            new_p3 = p3
+            new_p3 = line_cross_point(backward_edge, edge_opposite)
+            if point_dist_to_line(p0, p3, p1) > point_dist_to_line(p0, p3, p2):
+                # across p1
+                if backward_edge[1] == 0:
+                    backward_opposite = [1, 0, -p1[0]]
+                else:
+                    backward_opposite = [backward_edge[0], -1, p1[1] - backward_edge[0] * p1[0]]
+            else:
+                # across p2
+                if backward_edge[1] == 0:
+                    backward_opposite = [1, 0, -p2[0]]
+                else:
+                    backward_opposite = [backward_edge[0], -1, p2[1] - backward_edge[0] * p2[0]]
+            new_p1 = line_cross_point(backward_opposite, edge)
+            new_p2 = line_cross_point(backward_opposite, edge_opposite)
+            fitted_parallelograms.append([new_p0, new_p1, new_p2, new_p3, new_p0])
+
+        # 选定面积最小的平行四边形
+        areas = [Polygon(t).area for t in fitted_parallelograms]
+        parallelogram = np.array(fitted_parallelograms[np.argmin(areas)][:-1], dtype=np.float32)
+        # sort thie polygon
+        parallelogram_coord_sum = np.sum(parallelogram, axis=1)
+        min_coord_idx = np.argmin(parallelogram_coord_sum)
+        parallelogram = parallelogram[
+            [min_coord_idx, (min_coord_idx + 1) % 4, (min_coord_idx + 2) % 4, (min_coord_idx + 3) % 4]]
+
+        # 得到外包矩形即旋转角
+        rectange = rectangle_from_parallelogram(parallelogram)
+        rectange, rotate_angle = sort_rectangle(rectange)
+
+        p0_rect, p1_rect, p2_rect, p3_rect = rectange
+        # 对当前新加入的文本框区域像素点，根据其到矩形四边的距离修改geo_map
+        for y, x in xy_in_poly:
+            point = np.array([x, y], dtype=np.float32)
+            # top
+            geo_map[y, x, 0] = point_dist_to_line(p0_rect, p1_rect, point)
+            # right
+            geo_map[y, x, 1] = point_dist_to_line(p1_rect, p2_rect, point)
+            # down
+            geo_map[y, x, 2] = point_dist_to_line(p2_rect, p3_rect, point)
+            # left
+            geo_map[y, x, 3] = point_dist_to_line(p3_rect, p0_rect, point)
+            # angle
+            geo_map[y, x, 4] = rotate_angle
+    return score_map, geo_map, training_mask
+
+
+def generator(index,
+              input_size=512,
+              background_ratio=3. / 8,  # 纯背景样本比例
+              random_scale=np.array([0.5, 1, 2.0, 3.0]),  # 提取多尺度图片信息
+              image_list=None):
+    try:
+        im_fn = image_list[index]
+        im = cv2.imread(im_fn)
+        if im is None:
+            print("can't find image")
+            return None, None, None, None, None
+        h, w, _ = im.shape
+        # 所以要把gt去掉
+        txt_fn = im_fn.replace(os.path.basename(im_fn).split('.')[1], 'txt')
+        if not os.path.exists(txt_fn):
+            print('text file {} does not exists'.format(txt_fn))
+            return None, None, None, None, None
+        # 加载标注框信息
+        text_polys, text_tags = load_annoataion(txt_fn)
+
+        text_polys, text_tags = check_and_validate_polys(text_polys, text_tags, (h, w))
+
+        # random scale this image,随机选择一种尺度
+        rd_scale = np.random.choice(random_scale)
+        im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
+        text_polys *= rd_scale
+
+        # random crop a area from image，3/8���选中的概率，裁剪纯背景的图片
+        if np.random.rand() < background_ratio:
+            # crop background
+            im, text_polys, text_tags = crop_area(im, text_polys, text_tags, crop_background=True)
+            if text_polys.shape[0] > 0:
+                # print("cannot find background")
+                return None, None, None, None, None
+            # pad and resize image
+            new_h, new_w, _ = im.shape
+            max_h_w_i = np.max([new_h, new_w, input_size])
+            im_padded = np.zeros((max_h_w_i, max_h_w_i, 3), dtype=np.uint8)
+            im_padded[:new_h, :new_w, :] = im.copy()
+            # 将裁剪后图片扩充成512*512的图片
+            im = cv2.resize(im_padded, dsize=(input_size, input_size))
+            score_map = np.zeros((input_size, input_size), dtype=np.uint8)
+            geo_map_channels = 5 if cfg.geometry == 'RBOX' else 8
+            geo_map = np.zeros((input_size, input_size, geo_map_channels), dtype=np.float32)
+            training_mask = np.ones((input_size, input_size), dtype=np.uint8)
+        else:
+            # 5 / 8的选中的概率，裁剪含文本信息的图片
+            im, text_polys, text_tags = crop_area(im, text_polys, text_tags, crop_background=False)
+            if text_polys.shape[0] == 0:
+                # print("cannot find txt ground")
+                return None, None, None, None, None
+            h, w, _ = im.shape
+            # pad the image to the training input size or the longer side of image
+            new_h, new_w, _ = im.shape
+            max_h_w_i = np.max([new_h, new_w, input_size])
+            im_padded = np.zeros((max_h_w_i, max_h_w_i, 3), dtype=np.uint8)
+            im_padded[:new_h, :new_w, :] = im.copy()
+            im = im_padded
+            # resize the image to input size
+            # 填充，resize图像至设定尺寸
+            new_h, new_w, _ = im.shape
+            resize_h = input_size
+            resize_w = input_size
+            im = cv2.resize(im, dsize=(resize_w, resize_h))
+            # 将文本框坐标标签等比例修改
+            resize_ratio_3_x = resize_w / float(new_w)
+            resize_ratio_3_y = resize_h / float(new_h)
+            text_polys[:, :, 0] *= resize_ratio_3_x
+            text_polys[:, :, 1] *= resize_ratio_3_y
+            new_h, new_w, _ = im.shape
+            score_map, geo_map, training_mask = generate_rbox((new_h, new_w), text_polys, text_tags)
+
+        # 将一个样本的样本内容和标签信息append
+        images = im[:,:,::-1].astype(np.float32)
+        # 文件名加入列表
+        image_fns = im_fn
+        # 512*512取提取四分之一行列
+        score_maps = score_map[::4, ::4, np.newaxis].astype(np.float32)
+        geo_maps = geo_map[::4, ::4, :].astype(np.float32)
+        training_masks = training_mask[::4, ::4, np.newaxis].astype(np.float32)
+        # 符合一个样本之后输出
+        return images, image_fns, score_maps, geo_maps, training_masks
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+
+        # print("Exception is exist!")
+        return None, None, None, None, None

IndicPhotoOCR/detection/east_utils.py

@@ -0,0 +1,283 @@
+
+import torch
+import os
+from torch.nn import init
+import cv2
+import numpy as np
+import time
+import requests
+
+from IndicPhotoOCR.detection import east_config as cfg
+from IndicPhotoOCR.detection import east_preprossing as preprossing
+from IndicPhotoOCR.detection import east_locality_aware_nms as locality_aware_nms
+
+
+
+# Example usage:
+model_info = {
+    "east": {
+        "paths": [ cfg.checkpoint, cfg.pretrained_basemodel_path],
+        "urls" : ["https://github.com/anikde/STocr/releases/download/e0.1.0/epoch_990_checkpoint.pth.tar", "https://github.com/anikde/STocr/releases/download/e0.1.0/mobilenet_v2.pth.tar"]
+    },
+}
+
+class ModelManager:
+    def __init__(self):
+        # self.root_model_dir = "bharatOCR/detection/"
+        pass
+
+    def download_model(self, url, path):
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:  # Filter out keep-alive chunks
+                        f.write(chunk)
+            print(f"Downloaded: {path}")
+        else:
+            print(f"Failed to download from {url}")
+
+    def ensure_model(self, model_name):
+        model_paths = model_info[model_name]["paths"]  # Changed to handle multiple paths
+        urls = model_info[model_name]["urls"]          # Changed to handle multiple URLs
+
+
+        for model_path, url in zip(model_paths, urls):
+            # full_model_path = os.path.join(self.root_model_dir, model_path)
+
+            # Ensure the model path directory exists
+            os.makedirs(os.path.dirname(os.path.join(*cfg.pretrained_basemodel_path.split("/"))), exist_ok=True)
+
+            if not os.path.exists(model_path):
+                print(f"Model not found locally. Downloading {model_name} from {url}...")
+                self.download_model(url, model_path)
+            else:
+                print(f"Model already exists at {model_path}. No need to download.")
+
+
+
+# # Initialize ModelManager and ensure Hindi models are downloaded
+model_manager = ModelManager()
+model_manager.ensure_model("east")
+
+
+
+def init_weights(m_list, init_type=cfg.init_type, gain=0.02):
+    print("EAST <==> Prepare <==> Init Network'{}' <==> Begin".format(cfg.init_type))
+    # this will apply to each layer
+    for m in m_list:
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')  # good for relu
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print("EAST <==> Prepare <==> Init Network'{}' <==> Done".format(cfg.init_type))
+
+
+def Loading_checkpoint(model, optimizer, scheduler, filename='checkpoint.pth.tar'):
+    """[summary]
+    [description]
+    Arguments:
+        state {[type]} -- [description] a dict describe some params
+    Keyword Arguments:
+        filename {str} -- [description] (default: {'checkpoint.pth.tar'})
+    """
+    weightpath = os.path.abspath(cfg.checkpoint)
+    print("EAST <==> Prepare <==> Loading checkpoint '{}' <==> Begin".format(weightpath))
+    checkpoint = torch.load(weightpath)
+    start_epoch = checkpoint['epoch'] + 1
+    model.load_state_dict(checkpoint['state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer'])
+    scheduler.load_state_dict(checkpoint['scheduler'])
+    print("EAST <==> Prepare <==> Loading checkpoint '{}' <==> Done".format(weightpath))
+
+    return start_epoch
+
+
+def save_checkpoint(epoch, model, optimizer, scheduler, filename='checkpoint.pth.tar'):
+    """[summary]
+    [description]
+    Arguments:
+        state {[type]} -- [description] a dict describe some params
+    Keyword Arguments:
+        filename {str} -- [description] (default: {'checkpoint.pth.tar'})
+    """
+    print('EAST <==> Save weight - epoch {} <==> Begin'.format(epoch))
+    state = {
+        'epoch': epoch,
+        'state_dict': model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'scheduler': scheduler.state_dict()
+    }
+    weight_dir = cfg.save_model_path
+    if not os.path.exists(weight_dir):
+        os.mkdir(weight_dir)
+    filename = 'epoch_' + str(epoch) + '_checkpoint.pth.tar'
+    file_path = os.path.join(weight_dir, filename)
+    torch.save(state, file_path)
+    print('EAST <==> Save weight - epoch {} <==> Done'.format(epoch))
+
+
+class Regularization(torch.nn.Module):
+    def __init__(self, model, weight_decay, p=2):
+        super(Regularization, self).__init__()
+        if weight_decay < 0:
+            print("param weight_decay can not <0")
+            exit(0)
+        self.model = model
+        self.weight_decay = weight_decay
+        self.p = p
+        self.weight_list = self.get_weight(model)
+        # self.weight_info(self.weight_list)
+
+    def to(self, device):
+        self.device = device
+        super().to(device)
+        return self
+
+    def forward(self, model):
+        self.weight_list = self.get_weight(model)  # 获得最新的权重
+        reg_loss = self.regularization_loss(self.weight_list, self.weight_decay, p=self.p)
+        return reg_loss
+
+    def get_weight(self, model):
+        weight_list = []
+        for name, param in model.named_parameters():
+            if 'weight' in name:
+                weight = (name, param)
+                weight_list.append(weight)
+        return weight_list
+
+    def regularization_loss(self, weight_list, weight_decay, p=2):
+        reg_loss = 0
+        for name, w in weight_list:
+            l2_reg = torch.norm(w, p=p)
+            reg_loss = reg_loss + l2_reg
+
+        reg_loss = weight_decay * reg_loss
+        return reg_loss
+
+    def weight_info(self, weight_list):
+        print("---------------regularization weight---------------")
+        for name, w in weight_list:
+            print(name)
+        print("---------------------------------------------------")
+
+
+def resize_image(im, max_side_len=2400):
+    '''
+    resize image to a size multiple of 32 which is required by the network
+    :param im: the resized image
+    :param max_side_len: limit of max image size to avoid out of memory in gpu
+    :return: the resized image and the resize ratio
+    '''
+    h, w, _ = im.shape
+
+    resize_w = w
+    resize_h = h
+
+    # limit the max side
+    """
+    if max(resize_h, resize_w) > max_side_len:
+        ratio = float(max_side_len) / resize_h if resize_h > resize_w else float(max_side_len) / resize_w
+    else:
+        ratio = 1.
+
+    resize_h = int(resize_h * ratio)
+    resize_w = int(resize_w * ratio)
+    """
+
+    resize_h = resize_h if resize_h % 32 == 0 else (resize_h // 32 - 1) * 32
+    resize_w = resize_w if resize_w % 32 == 0 else (resize_w // 32 - 1) * 32
+    #resize_h, resize_w = 512, 512
+    im = cv2.resize(im, (int(resize_w), int(resize_h)))
+
+    ratio_h = resize_h / float(h)
+    ratio_w = resize_w / float(w)
+
+    return im, (ratio_h, ratio_w)
+
+
+def detect(score_map, geo_map, timer, score_map_thresh=0.8, box_thresh=0.1, nms_thres=0.2):
+    '''
+    restore text boxes from score map and geo map
+    :param score_map:
+    :param geo_map:
+    :param timer:
+    :param score_map_thresh: threshhold for score map
+    :param box_thresh: threshhold for boxes
+    :param nms_thres: threshold for nms
+    :return:
+    '''
+
+    # score_map 和 geo_map 的维数进行调整
+    if len(score_map.shape) == 4:
+        score_map = score_map[0, :, :, 0]
+        geo_map = geo_map[0, :, :, :]
+    # filter the score map
+    xy_text = np.argwhere(score_map > score_map_thresh)
+    # sort the text boxes via the y axis
+    xy_text = xy_text[np.argsort(xy_text[:, 0])]
+    # restore
+    start = time.time()
+    text_box_restored = preprossing.restore_rectangle(xy_text[:, ::-1] * 4,
+                                                      geo_map[xy_text[:, 0], xy_text[:, 1], :])  # N*4*2
+    # print('{} text boxes before nms'.format(text_box_restored.shape[0]))
+    boxes = np.zeros((text_box_restored.shape[0], 9), dtype=np.float32)
+    boxes[:, :8] = text_box_restored.reshape((-1, 8))
+    boxes[:, 8] = score_map[xy_text[:, 0], xy_text[:, 1]]
+    timer['restore'] = time.time() - start
+    # nms part
+    start = time.time()
+    boxes = locality_aware_nms.nms_locality(boxes.astype(np.float64), nms_thres)
+    timer['nms'] = time.time() - start
+    # print(timer['nms'])
+    if boxes.shape[0] == 0:
+        return None, timer
+
+    # here we filter some low score boxes by the average score map, this is different from the orginal paper
+    for i, box in enumerate(boxes):
+        mask = np.zeros_like(score_map, dtype=np.uint8)
+        cv2.fillPoly(mask, box[:8].reshape((-1, 4, 2)).astype(np.int32) // 4, 1)
+        boxes[i, 8] = cv2.mean(score_map, mask)[0]
+    boxes = boxes[boxes[:, 8] > box_thresh]
+    return boxes, timer
+
+
+def sort_poly(p):
+    min_axis = np.argmin(np.sum(p, axis=1))
+    p = p[[min_axis, (min_axis + 1) % 4, (min_axis + 2) % 4, (min_axis + 3) % 4]]
+    if abs(p[0, 0] - p[1, 0]) > abs(p[0, 1] - p[1, 1]):
+        return p
+    else:
+        return p[[0, 3, 2, 1]]
+
+
+def mean_image_subtraction(images, means=cfg.means):
+    '''
+    image normalization
+    :param images: bs * w * h * channel
+    :param means:
+    :return:
+    '''
+    num_channels = images.data.shape[1]
+    if len(means) != num_channels:
+        raise ValueError('len(means) must match the number of channels')
+    for i in range(num_channels):
+        images.data[:, i, :, :] -= means[i]
+
+    return images

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())