# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function import paddle from paddle import nn import numpy as np import cv2 __all__ = ["Kie_backbone"] class Encoder(nn.Layer): def __init__(self, num_channels, num_filters): super(Encoder, self).__init__() self.conv1 = nn.Conv2D( num_channels, num_filters, kernel_size=3, stride=1, padding=1, bias_attr=False) self.bn1 = nn.BatchNorm(num_filters, act='relu') self.conv2 = nn.Conv2D( num_filters, num_filters, kernel_size=3, stride=1, padding=1, bias_attr=False) self.bn2 = nn.BatchNorm(num_filters, act='relu') self.pool = nn.MaxPool2D(kernel_size=3, stride=2, padding=1) def forward(self, inputs): x = self.conv1(inputs) x = self.bn1(x) x = self.conv2(x) x = self.bn2(x) x_pooled = self.pool(x) return x, x_pooled class Decoder(nn.Layer): def __init__(self, num_channels, num_filters): super(Decoder, self).__init__() self.conv1 = nn.Conv2D( num_channels, num_filters, kernel_size=3, stride=1, padding=1, bias_attr=False) self.bn1 = nn.BatchNorm(num_filters, act='relu') self.conv2 = nn.Conv2D( num_filters, num_filters, kernel_size=3, stride=1, padding=1, bias_attr=False) self.bn2 = nn.BatchNorm(num_filters, act='relu') self.conv0 = nn.Conv2D( num_channels, num_filters, kernel_size=1, stride=1, padding=0, bias_attr=False) self.bn0 = nn.BatchNorm(num_filters, act='relu') def forward(self, inputs_prev, inputs): x = self.conv0(inputs) x = self.bn0(x) x = paddle.nn.functional.interpolate( x, scale_factor=2, mode='bilinear', align_corners=False) x = paddle.concat([inputs_prev, x], axis=1) x = self.conv1(x) x = self.bn1(x) x = self.conv2(x) x = self.bn2(x) return x class UNet(nn.Layer): def __init__(self): super(UNet, self).__init__() self.down1 = Encoder(num_channels=3, num_filters=16) self.down2 = Encoder(num_channels=16, num_filters=32) self.down3 = Encoder(num_channels=32, num_filters=64) self.down4 = Encoder(num_channels=64, num_filters=128) self.down5 = Encoder(num_channels=128, num_filters=256) self.up1 = Decoder(32, 16) self.up2 = Decoder(64, 32) self.up3 = Decoder(128, 64) self.up4 = Decoder(256, 128) self.out_channels = 16 def forward(self, inputs): x1, _ = self.down1(inputs) _, x2 = self.down2(x1) _, x3 = self.down3(x2) _, x4 = self.down4(x3) _, x5 = self.down5(x4) x = self.up4(x4, x5) x = self.up3(x3, x) x = self.up2(x2, x) x = self.up1(x1, x) return x class Kie_backbone(nn.Layer): def __init__(self, in_channels, **kwargs): super(Kie_backbone, self).__init__() self.out_channels = 16 self.img_feat = UNet() self.maxpool = nn.MaxPool2D(kernel_size=7) def bbox2roi(self, bbox_list): rois_list = [] rois_num = [] for img_id, bboxes in enumerate(bbox_list): rois_num.append(bboxes.shape[0]) rois_list.append(bboxes) rois = paddle.concat(rois_list, 0) rois_num = paddle.to_tensor(rois_num, dtype='int32') return rois, rois_num def pre_process(self, img, relations, texts, gt_bboxes, tag, img_size): img, relations, texts, gt_bboxes, tag, img_size = img.numpy( ), relations.numpy(), texts.numpy(), gt_bboxes.numpy(), tag.numpy( ).tolist(), img_size.numpy() temp_relations, temp_texts, temp_gt_bboxes = [], [], [] h, w = int(np.max(img_size[:, 0])), int(np.max(img_size[:, 1])) img = paddle.to_tensor(img[:, :, :h, :w]) batch = len(tag) for i in range(batch): num, recoder_len = tag[i][0], tag[i][1] temp_relations.append( paddle.to_tensor( relations[i, :num, :num, :], dtype='float32')) temp_texts.append( paddle.to_tensor( texts[i, :num, :recoder_len], dtype='float32')) temp_gt_bboxes.append( paddle.to_tensor( gt_bboxes[i, :num, ...], dtype='float32')) return img, temp_relations, temp_texts, temp_gt_bboxes def forward(self, inputs): img, relations, texts, gt_bboxes, tag, img_size = inputs[0], inputs[ 1], inputs[2], inputs[3], inputs[5], inputs[-1] img, relations, texts, gt_bboxes = self.pre_process( img, relations, texts, gt_bboxes, tag, img_size) # for i in range(4): # img_t = (img[i].numpy().transpose([1, 2, 0]) * 255.0).astype('uint8') # img_t = img_t.copy() # gt_bboxes_t = gt_bboxes[i].cpu().numpy() # box = gt_bboxes_t.astype(np.int32).reshape((-1, 1, 2)) # cv2.polylines(img_t, [box], True, color=(255, 255, 0), thickness=1) # cv2.imwrite("/Users/hongyongjie/project/PaddleOCR/output/{}.png".format(i), img_t) # # cv2.imwrite("/Users/hongyongjie/project/PaddleOCR/output/{}.png".format(i), img_t * 255.0) # exit() x = self.img_feat(img) boxes, rois_num = self.bbox2roi(gt_bboxes) feats = paddle.fluid.layers.roi_align( x, boxes, spatial_scale=1.0, pooled_height=7, pooled_width=7, rois_num=rois_num) feats = self.maxpool(feats).squeeze(-1).squeeze(-1) return [relations, texts, feats]