craft-detector 0.2.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- craft_detector/__init__.py +3 -0
- craft_detector/basenet/__init__.py +0 -0
- craft_detector/basenet/vgg16_bn.py +75 -0
- craft_detector/craft.py +83 -0
- craft_detector/craft_utils.py +328 -0
- craft_detector/file_utils.py +76 -0
- craft_detector/imgproc.py +68 -0
- craft_detector/refinenet.py +65 -0
- craft_detector/utils/logging.py +21 -0
- craft_detector/utils/utils.py +264 -0
- craft_detector-0.2.1.dist-info/METADATA +130 -0
- craft_detector-0.2.1.dist-info/RECORD +14 -0
- craft_detector-0.2.1.dist-info/WHEEL +4 -0
- craft_detector-0.2.1.dist-info/licenses/LICENSE +19 -0
|
File without changes
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
from collections import namedtuple
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
import torch.nn as nn
|
|
5
|
+
import torch.nn.init as init
|
|
6
|
+
from torchvision import models
|
|
7
|
+
from torchvision.models.vgg import VGG16_BN_Weights
|
|
8
|
+
|
|
9
|
+
def init_weights(modules):
|
|
10
|
+
for m in modules:
|
|
11
|
+
if isinstance(m, nn.Conv2d):
|
|
12
|
+
init.xavier_uniform_(m.weight.data)
|
|
13
|
+
if m.bias is not None:
|
|
14
|
+
m.bias.data.zero_()
|
|
15
|
+
elif isinstance(m, nn.BatchNorm2d):
|
|
16
|
+
m.weight.data.fill_(1)
|
|
17
|
+
m.bias.data.zero_()
|
|
18
|
+
elif isinstance(m, nn.Linear):
|
|
19
|
+
m.weight.data.normal_(0, 0.01)
|
|
20
|
+
m.bias.data.zero_()
|
|
21
|
+
|
|
22
|
+
class vgg16_bn(torch.nn.Module):
|
|
23
|
+
def __init__(self, pretrained=True, freeze=True):
|
|
24
|
+
super(vgg16_bn, self).__init__()
|
|
25
|
+
# Use the weights parameter based on the pretrained flag
|
|
26
|
+
weights = VGG16_BN_Weights.IMAGENET1K_V1 if pretrained else None
|
|
27
|
+
vgg_pretrained_features = models.vgg16_bn(weights=weights).features
|
|
28
|
+
|
|
29
|
+
self.slice1 = torch.nn.Sequential()
|
|
30
|
+
self.slice2 = torch.nn.Sequential()
|
|
31
|
+
self.slice3 = torch.nn.Sequential()
|
|
32
|
+
self.slice4 = torch.nn.Sequential()
|
|
33
|
+
self.slice5 = torch.nn.Sequential()
|
|
34
|
+
for x in range(12): # conv2_2
|
|
35
|
+
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
|
36
|
+
for x in range(12, 19): # conv3_3
|
|
37
|
+
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
|
38
|
+
for x in range(19, 29): # conv4_3
|
|
39
|
+
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
|
40
|
+
for x in range(29, 39): # conv5_3
|
|
41
|
+
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
|
42
|
+
|
|
43
|
+
# fc6, fc7 without atrous conv
|
|
44
|
+
self.slice5 = torch.nn.Sequential(
|
|
45
|
+
nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
|
|
46
|
+
nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6),
|
|
47
|
+
nn.Conv2d(1024, 1024, kernel_size=1)
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
if not pretrained:
|
|
51
|
+
init_weights(self.slice1.modules())
|
|
52
|
+
init_weights(self.slice2.modules())
|
|
53
|
+
init_weights(self.slice3.modules())
|
|
54
|
+
init_weights(self.slice4.modules())
|
|
55
|
+
|
|
56
|
+
init_weights(self.slice5.modules()) # no pretrained model for fc6 and fc7
|
|
57
|
+
|
|
58
|
+
if freeze:
|
|
59
|
+
for param in self.slice1.parameters(): # only first conv
|
|
60
|
+
param.requires_grad= False
|
|
61
|
+
|
|
62
|
+
def forward(self, X):
|
|
63
|
+
h = self.slice1(X)
|
|
64
|
+
h_relu2_2 = h
|
|
65
|
+
h = self.slice2(h)
|
|
66
|
+
h_relu3_2 = h
|
|
67
|
+
h = self.slice3(h)
|
|
68
|
+
h_relu4_3 = h
|
|
69
|
+
h = self.slice4(h)
|
|
70
|
+
h_relu5_3 = h
|
|
71
|
+
h = self.slice5(h)
|
|
72
|
+
h_fc7 = h
|
|
73
|
+
vgg_outputs = namedtuple("VggOutputs", ['fc7', 'relu5_3', 'relu4_3', 'relu3_2', 'relu2_2'])
|
|
74
|
+
out = vgg_outputs(h_fc7, h_relu5_3, h_relu4_3, h_relu3_2, h_relu2_2)
|
|
75
|
+
return out
|
craft_detector/craft.py
ADDED
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
3
|
+
MIT License
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
# -*- coding: utf-8 -*-
|
|
7
|
+
import torch
|
|
8
|
+
import torch.nn as nn
|
|
9
|
+
import torch.nn.functional as F
|
|
10
|
+
|
|
11
|
+
from .basenet.vgg16_bn import vgg16_bn, init_weights
|
|
12
|
+
|
|
13
|
+
class double_conv(nn.Module):
|
|
14
|
+
def __init__(self, in_ch, mid_ch, out_ch):
|
|
15
|
+
super(double_conv, self).__init__()
|
|
16
|
+
self.conv = nn.Sequential(
|
|
17
|
+
nn.Conv2d(in_ch + mid_ch, mid_ch, kernel_size=1),
|
|
18
|
+
nn.BatchNorm2d(mid_ch),
|
|
19
|
+
nn.ReLU(inplace=True),
|
|
20
|
+
nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1),
|
|
21
|
+
nn.BatchNorm2d(out_ch),
|
|
22
|
+
nn.ReLU(inplace=True)
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
def forward(self, x):
|
|
26
|
+
x = self.conv(x)
|
|
27
|
+
return x
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class CRAFT(nn.Module):
|
|
31
|
+
def __init__(self, pretrained=False, freeze=False):
|
|
32
|
+
super(CRAFT, self).__init__()
|
|
33
|
+
|
|
34
|
+
""" Base network """
|
|
35
|
+
self.basenet = vgg16_bn(pretrained, freeze)
|
|
36
|
+
|
|
37
|
+
""" U network """
|
|
38
|
+
self.upconv1 = double_conv(1024, 512, 256)
|
|
39
|
+
self.upconv2 = double_conv(512, 256, 128)
|
|
40
|
+
self.upconv3 = double_conv(256, 128, 64)
|
|
41
|
+
self.upconv4 = double_conv(128, 64, 32)
|
|
42
|
+
|
|
43
|
+
num_class = 2
|
|
44
|
+
self.conv_cls = nn.Sequential(
|
|
45
|
+
nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
|
|
46
|
+
nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
|
|
47
|
+
nn.Conv2d(32, 16, kernel_size=3, padding=1), nn.ReLU(inplace=True),
|
|
48
|
+
nn.Conv2d(16, 16, kernel_size=1), nn.ReLU(inplace=True),
|
|
49
|
+
nn.Conv2d(16, num_class, kernel_size=1),
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
init_weights(self.upconv1.modules())
|
|
53
|
+
init_weights(self.upconv2.modules())
|
|
54
|
+
init_weights(self.upconv3.modules())
|
|
55
|
+
init_weights(self.upconv4.modules())
|
|
56
|
+
init_weights(self.conv_cls.modules())
|
|
57
|
+
|
|
58
|
+
def _forward_impl(self, x):
|
|
59
|
+
""" Base network """
|
|
60
|
+
sources = self.basenet(x)
|
|
61
|
+
|
|
62
|
+
""" U network """
|
|
63
|
+
y = torch.cat([sources[0], sources[1]], dim=1)
|
|
64
|
+
y = self.upconv1(y)
|
|
65
|
+
|
|
66
|
+
y = F.interpolate(y, size=sources[2].size()[2:], mode='bilinear', align_corners=False)
|
|
67
|
+
y = torch.cat([y, sources[2]], dim=1)
|
|
68
|
+
y = self.upconv2(y)
|
|
69
|
+
|
|
70
|
+
y = F.interpolate(y, size=sources[3].size()[2:], mode='bilinear', align_corners=False)
|
|
71
|
+
y = torch.cat([y, sources[3]], dim=1)
|
|
72
|
+
y = self.upconv3(y)
|
|
73
|
+
|
|
74
|
+
y = F.interpolate(y, size=sources[4].size()[2:], mode='bilinear', align_corners=False)
|
|
75
|
+
y = torch.cat([y, sources[4]], dim=1)
|
|
76
|
+
feature = self.upconv4(y)
|
|
77
|
+
|
|
78
|
+
y = self.conv_cls(feature)
|
|
79
|
+
|
|
80
|
+
return y.permute(0,2,3,1), feature
|
|
81
|
+
|
|
82
|
+
def forward(self, x):
|
|
83
|
+
return self._forward_impl(x)
|
|
@@ -0,0 +1,328 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
3
|
+
MIT License
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
# -*- coding: utf-8 -*-
|
|
7
|
+
import numpy as np
|
|
8
|
+
import cv2
|
|
9
|
+
import math
|
|
10
|
+
import torch
|
|
11
|
+
|
|
12
|
+
from scipy.ndimage import label
|
|
13
|
+
from torch.autograd import Variable
|
|
14
|
+
|
|
15
|
+
from .imgproc import normalizeMeanVariance
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
""" auxiliary functions """
|
|
19
|
+
# unwarp corodinates
|
|
20
|
+
def warpCoord(Minv, pt):
|
|
21
|
+
out = np.matmul(Minv, (pt[0], pt[1], 1))
|
|
22
|
+
return np.array([out[0]/out[2], out[1]/out[2]])
|
|
23
|
+
""" end of auxiliary functions """
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text, estimate_num_chars=False):
|
|
27
|
+
# prepare data
|
|
28
|
+
linkmap = linkmap.copy()
|
|
29
|
+
textmap = textmap.copy()
|
|
30
|
+
img_h, img_w = textmap.shape
|
|
31
|
+
|
|
32
|
+
""" labeling method """
|
|
33
|
+
ret, text_score = cv2.threshold(textmap, low_text, 1, 0)
|
|
34
|
+
ret, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)
|
|
35
|
+
|
|
36
|
+
text_score_comb = np.clip(text_score + link_score, 0, 1)
|
|
37
|
+
nLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(text_score_comb.astype(np.uint8), connectivity=4)
|
|
38
|
+
|
|
39
|
+
"""
|
|
40
|
+
print("nlabels = ", nLabels)
|
|
41
|
+
xx = labels == 1
|
|
42
|
+
print("xx.shape = ", xx.shape)
|
|
43
|
+
print("textmap[xx].shape = ", textmap[xx].shape)
|
|
44
|
+
print("np.max(textmap[xx]) = ", np.max(textmap[xx]))
|
|
45
|
+
size = stats[1, cv2.CC_STAT_AREA]
|
|
46
|
+
print("area =", size)
|
|
47
|
+
exit(0)
|
|
48
|
+
"""
|
|
49
|
+
|
|
50
|
+
det = []
|
|
51
|
+
mapper = []
|
|
52
|
+
for k in range(1,nLabels):
|
|
53
|
+
# size filtering
|
|
54
|
+
size = stats[k, cv2.CC_STAT_AREA]
|
|
55
|
+
#if size < 10: continue
|
|
56
|
+
|
|
57
|
+
# thresholding
|
|
58
|
+
if np.max(textmap[labels==k]) < text_threshold: continue
|
|
59
|
+
|
|
60
|
+
# make segmentation map
|
|
61
|
+
segmap = np.zeros(textmap.shape, dtype=np.uint8)
|
|
62
|
+
segmap[labels==k] = 255
|
|
63
|
+
if estimate_num_chars:
|
|
64
|
+
_, character_locs = cv2.threshold((textmap - linkmap) * segmap /255., text_threshold, 1, 0)
|
|
65
|
+
_, n_chars = label(character_locs)
|
|
66
|
+
mapper.append(n_chars)
|
|
67
|
+
else:
|
|
68
|
+
mapper.append(k)
|
|
69
|
+
segmap[np.logical_and(link_score==1, text_score==0)] = 0 # remove link area
|
|
70
|
+
x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
|
|
71
|
+
w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
|
|
72
|
+
niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
|
|
73
|
+
sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
|
|
74
|
+
# boundary check
|
|
75
|
+
if sx < 0 : sx = 0
|
|
76
|
+
if sy < 0 : sy = 0
|
|
77
|
+
if ex >= img_w: ex = img_w
|
|
78
|
+
if ey >= img_h: ey = img_h
|
|
79
|
+
kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(1 + niter, 1 + niter))
|
|
80
|
+
segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)
|
|
81
|
+
|
|
82
|
+
# make box
|
|
83
|
+
np_contours = np.roll(np.array(np.where(segmap!=0)),1,axis=0).transpose().reshape(-1,2)
|
|
84
|
+
rectangle = cv2.minAreaRect(np_contours)
|
|
85
|
+
box = cv2.boxPoints(rectangle)
|
|
86
|
+
|
|
87
|
+
# align diamond-shape
|
|
88
|
+
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
|
|
89
|
+
box_ratio = max(w, h) / (min(w, h) + 1e-5)
|
|
90
|
+
if abs(1 - box_ratio) <= 0.1:
|
|
91
|
+
l, r = min(np_contours[:,0]), max(np_contours[:,0])
|
|
92
|
+
t, b = min(np_contours[:,1]), max(np_contours[:,1])
|
|
93
|
+
box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
|
|
94
|
+
|
|
95
|
+
# make clock-wise order
|
|
96
|
+
startidx = box.sum(axis=1).argmin()
|
|
97
|
+
box = np.roll(box, 4-startidx, 0)
|
|
98
|
+
box = np.array(box)
|
|
99
|
+
|
|
100
|
+
det.append(box)
|
|
101
|
+
|
|
102
|
+
return det, labels, mapper
|
|
103
|
+
|
|
104
|
+
def getPoly_core(boxes, labels, mapper, linkmap):
|
|
105
|
+
# configs
|
|
106
|
+
num_cp = 5
|
|
107
|
+
max_len_ratio = 0.7
|
|
108
|
+
expand_ratio = 1.45
|
|
109
|
+
max_r = 2.0
|
|
110
|
+
step_r = 0.2
|
|
111
|
+
|
|
112
|
+
polys = []
|
|
113
|
+
for k, box in enumerate(boxes):
|
|
114
|
+
# size filter for small instance
|
|
115
|
+
w, h = int(np.linalg.norm(box[0] - box[1]) + 1), int(np.linalg.norm(box[1] - box[2]) + 1)
|
|
116
|
+
if w < 10 or h < 10:
|
|
117
|
+
polys.append(None); continue
|
|
118
|
+
|
|
119
|
+
# warp image
|
|
120
|
+
tar = np.float32([[0,0],[w,0],[w,h],[0,h]])
|
|
121
|
+
M = cv2.getPerspectiveTransform(box, tar)
|
|
122
|
+
word_label = cv2.warpPerspective(labels, M, (w, h), flags=cv2.INTER_NEAREST)
|
|
123
|
+
try:
|
|
124
|
+
Minv = np.linalg.inv(M)
|
|
125
|
+
except:
|
|
126
|
+
polys.append(None); continue
|
|
127
|
+
|
|
128
|
+
# binarization for selected label
|
|
129
|
+
cur_label = mapper[k]
|
|
130
|
+
word_label[word_label != cur_label] = 0
|
|
131
|
+
word_label[word_label > 0] = 1
|
|
132
|
+
|
|
133
|
+
""" Polygon generation """
|
|
134
|
+
# find top/bottom contours
|
|
135
|
+
cp = []
|
|
136
|
+
max_len = -1
|
|
137
|
+
for i in range(w):
|
|
138
|
+
region = np.where(word_label[:,i] != 0)[0]
|
|
139
|
+
if len(region) < 2 : continue
|
|
140
|
+
cp.append((i, region[0], region[-1]))
|
|
141
|
+
length = region[-1] - region[0] + 1
|
|
142
|
+
if length > max_len: max_len = length
|
|
143
|
+
|
|
144
|
+
# pass if max_len is similar to h
|
|
145
|
+
if h * max_len_ratio < max_len:
|
|
146
|
+
polys.append(None); continue
|
|
147
|
+
|
|
148
|
+
# get pivot points with fixed length
|
|
149
|
+
tot_seg = num_cp * 2 + 1
|
|
150
|
+
seg_w = w / tot_seg # segment width
|
|
151
|
+
pp = [None] * num_cp # init pivot points
|
|
152
|
+
cp_section = [[0, 0]] * tot_seg
|
|
153
|
+
seg_height = [0] * num_cp
|
|
154
|
+
seg_num = 0
|
|
155
|
+
num_sec = 0
|
|
156
|
+
prev_h = -1
|
|
157
|
+
for i in range(0,len(cp)):
|
|
158
|
+
(x, sy, ey) = cp[i]
|
|
159
|
+
if (seg_num + 1) * seg_w <= x and seg_num <= tot_seg:
|
|
160
|
+
# average previous segment
|
|
161
|
+
if num_sec == 0: break
|
|
162
|
+
cp_section[seg_num] = [cp_section[seg_num][0] / num_sec, cp_section[seg_num][1] / num_sec]
|
|
163
|
+
num_sec = 0
|
|
164
|
+
|
|
165
|
+
# reset variables
|
|
166
|
+
seg_num += 1
|
|
167
|
+
prev_h = -1
|
|
168
|
+
|
|
169
|
+
# accumulate center points
|
|
170
|
+
cy = (sy + ey) * 0.5
|
|
171
|
+
cur_h = ey - sy + 1
|
|
172
|
+
cp_section[seg_num] = [cp_section[seg_num][0] + x, cp_section[seg_num][1] + cy]
|
|
173
|
+
num_sec += 1
|
|
174
|
+
|
|
175
|
+
if seg_num % 2 == 0: continue # No polygon area
|
|
176
|
+
|
|
177
|
+
if prev_h < cur_h:
|
|
178
|
+
pp[int((seg_num - 1)/2)] = (x, cy)
|
|
179
|
+
seg_height[int((seg_num - 1)/2)] = cur_h
|
|
180
|
+
prev_h = cur_h
|
|
181
|
+
|
|
182
|
+
# processing last segment
|
|
183
|
+
if num_sec != 0:
|
|
184
|
+
cp_section[-1] = [cp_section[-1][0] / num_sec, cp_section[-1][1] / num_sec]
|
|
185
|
+
|
|
186
|
+
# pass if num of pivots is not sufficient or segment width is smaller than character height
|
|
187
|
+
if None in pp or seg_w < np.max(seg_height) * 0.25:
|
|
188
|
+
polys.append(None); continue
|
|
189
|
+
|
|
190
|
+
# calc median maximum of pivot points
|
|
191
|
+
half_char_h = np.median(seg_height) * expand_ratio / 2
|
|
192
|
+
|
|
193
|
+
# calc gradiant and apply to make horizontal pivots
|
|
194
|
+
new_pp = []
|
|
195
|
+
for i, (x, cy) in enumerate(pp):
|
|
196
|
+
dx = cp_section[i * 2 + 2][0] - cp_section[i * 2][0]
|
|
197
|
+
dy = cp_section[i * 2 + 2][1] - cp_section[i * 2][1]
|
|
198
|
+
if dx == 0: # gradient if zero
|
|
199
|
+
new_pp.append([x, cy - half_char_h, x, cy + half_char_h])
|
|
200
|
+
continue
|
|
201
|
+
rad = - math.atan2(dy, dx)
|
|
202
|
+
c, s = half_char_h * math.cos(rad), half_char_h * math.sin(rad)
|
|
203
|
+
new_pp.append([x - s, cy - c, x + s, cy + c])
|
|
204
|
+
|
|
205
|
+
# get edge points to cover character heatmaps
|
|
206
|
+
isSppFound, isEppFound = False, False
|
|
207
|
+
grad_s = (pp[1][1] - pp[0][1]) / (pp[1][0] - pp[0][0]) + (pp[2][1] - pp[1][1]) / (pp[2][0] - pp[1][0])
|
|
208
|
+
grad_e = (pp[-2][1] - pp[-1][1]) / (pp[-2][0] - pp[-1][0]) + (pp[-3][1] - pp[-2][1]) / (pp[-3][0] - pp[-2][0])
|
|
209
|
+
for r in np.arange(0.5, max_r, step_r):
|
|
210
|
+
dx = 2 * half_char_h * r
|
|
211
|
+
if not isSppFound:
|
|
212
|
+
line_img = np.zeros(word_label.shape, dtype=np.uint8)
|
|
213
|
+
dy = grad_s * dx
|
|
214
|
+
p = np.array(new_pp[0]) - np.array([dx, dy, dx, dy])
|
|
215
|
+
cv2.line(line_img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), 1, thickness=1)
|
|
216
|
+
if np.sum(np.logical_and(word_label, line_img)) == 0 or r + 2 * step_r >= max_r:
|
|
217
|
+
spp = p
|
|
218
|
+
isSppFound = True
|
|
219
|
+
if not isEppFound:
|
|
220
|
+
line_img = np.zeros(word_label.shape, dtype=np.uint8)
|
|
221
|
+
dy = grad_e * dx
|
|
222
|
+
p = np.array(new_pp[-1]) + np.array([dx, dy, dx, dy])
|
|
223
|
+
cv2.line(line_img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), 1, thickness=1)
|
|
224
|
+
if np.sum(np.logical_and(word_label, line_img)) == 0 or r + 2 * step_r >= max_r:
|
|
225
|
+
epp = p
|
|
226
|
+
isEppFound = True
|
|
227
|
+
if isSppFound and isEppFound:
|
|
228
|
+
break
|
|
229
|
+
|
|
230
|
+
# pass if boundary of polygon is not found
|
|
231
|
+
if not (isSppFound and isEppFound):
|
|
232
|
+
polys.append(None); continue
|
|
233
|
+
|
|
234
|
+
# make final polygon
|
|
235
|
+
poly = []
|
|
236
|
+
poly.append(warpCoord(Minv, (spp[0], spp[1])))
|
|
237
|
+
for p in new_pp:
|
|
238
|
+
poly.append(warpCoord(Minv, (p[0], p[1])))
|
|
239
|
+
poly.append(warpCoord(Minv, (epp[0], epp[1])))
|
|
240
|
+
poly.append(warpCoord(Minv, (epp[2], epp[3])))
|
|
241
|
+
for p in reversed(new_pp):
|
|
242
|
+
poly.append(warpCoord(Minv, (p[2], p[3])))
|
|
243
|
+
poly.append(warpCoord(Minv, (spp[2], spp[3])))
|
|
244
|
+
|
|
245
|
+
# add to final result
|
|
246
|
+
polys.append(np.array(poly))
|
|
247
|
+
|
|
248
|
+
return polys
|
|
249
|
+
|
|
250
|
+
def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False, estimate_num_chars=False):
|
|
251
|
+
if poly and estimate_num_chars:
|
|
252
|
+
raise Exception("Estimating the number of characters not currently supported with poly.")
|
|
253
|
+
boxes, labels, mapper = getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text, estimate_num_chars)
|
|
254
|
+
|
|
255
|
+
if poly:
|
|
256
|
+
polys = getPoly_core(boxes, labels, mapper, linkmap)
|
|
257
|
+
else:
|
|
258
|
+
polys = [None] * len(boxes)
|
|
259
|
+
|
|
260
|
+
return boxes, polys, mapper
|
|
261
|
+
|
|
262
|
+
def adjustResultCoordinates(polys, ratio_w, ratio_h, ratio_net = 2):
|
|
263
|
+
if len(polys) > 0:
|
|
264
|
+
for k in range(len(polys)):
|
|
265
|
+
if polys[k] is not None:
|
|
266
|
+
polys[k] *= (ratio_w * ratio_net, ratio_h * ratio_net)
|
|
267
|
+
return polys
|
|
268
|
+
|
|
269
|
+
|
|
270
|
+
def detect_net(image, net, refine_net=None, device=None, half=False):
|
|
271
|
+
|
|
272
|
+
x = None
|
|
273
|
+
if isinstance(image, np.ndarray):
|
|
274
|
+
if image.dtype != np.float32:
|
|
275
|
+
# for normal ndarray images
|
|
276
|
+
x = normalizeMeanVariance(image)
|
|
277
|
+
if x is None:
|
|
278
|
+
x = image.copy()
|
|
279
|
+
|
|
280
|
+
# preprocessing
|
|
281
|
+
x = torch.from_numpy(x).permute(2, 0, 1) # [h, w, c] to [c, h, w]
|
|
282
|
+
x = Variable(x.unsqueeze(0)) # [c, h, w] to [b, c, h, w]
|
|
283
|
+
|
|
284
|
+
if device:
|
|
285
|
+
x = x.to(device)
|
|
286
|
+
|
|
287
|
+
if half:
|
|
288
|
+
x = x.half()
|
|
289
|
+
|
|
290
|
+
# forward pass
|
|
291
|
+
with torch.no_grad():
|
|
292
|
+
y, feature = net(x)
|
|
293
|
+
|
|
294
|
+
del x
|
|
295
|
+
|
|
296
|
+
# make score and link map
|
|
297
|
+
if half:
|
|
298
|
+
score_text = y[0, ..., 0].to(torch.float32).detach().cpu().numpy()
|
|
299
|
+
score_link = y[0, ..., 1].to(torch.float32).detach().cpu().numpy()
|
|
300
|
+
else:
|
|
301
|
+
score_text = y[0, ..., 0].detach().cpu().numpy()
|
|
302
|
+
score_link = y[0, ..., 1].detach().cpu().numpy()
|
|
303
|
+
|
|
304
|
+
# refine link
|
|
305
|
+
if refine_net is not None:
|
|
306
|
+
with torch.no_grad():
|
|
307
|
+
y_refiner = refine_net(y, feature)
|
|
308
|
+
|
|
309
|
+
if half:
|
|
310
|
+
score_link = y_refiner[0, ...,0].to(torch.float32).detach().cpu().numpy()
|
|
311
|
+
else:
|
|
312
|
+
score_link = y_refiner[0, ...,0].detach().cpu().numpy()
|
|
313
|
+
|
|
314
|
+
del y, feature
|
|
315
|
+
|
|
316
|
+
return score_text, score_link
|
|
317
|
+
|
|
318
|
+
|
|
319
|
+
def detect_textbox(image, net, refine_net, text_threshold, link_threshold, low_text, poly, device, estimate_num_chars=False,
|
|
320
|
+
half=False):
|
|
321
|
+
|
|
322
|
+
score_text, score_link = detect_net(image, net, refine_net, device=device, half=half)
|
|
323
|
+
|
|
324
|
+
# Post-processing
|
|
325
|
+
boxes, polys, mapper = getDetBoxes(score_text, score_link, text_threshold, link_threshold, low_text, poly,
|
|
326
|
+
estimate_num_chars=estimate_num_chars)
|
|
327
|
+
|
|
328
|
+
return score_text, score_link, boxes, polys
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
# -*- coding: utf-8 -*-
|
|
2
|
+
import os
|
|
3
|
+
import numpy as np
|
|
4
|
+
import cv2
|
|
5
|
+
from . import imgproc
|
|
6
|
+
|
|
7
|
+
# borrowed from https://github.com/lengstrom/fast-style-transfer/blob/master/src/utils.py
|
|
8
|
+
def get_files(img_dir):
|
|
9
|
+
imgs, masks, xmls = list_files(img_dir)
|
|
10
|
+
return imgs, masks, xmls
|
|
11
|
+
|
|
12
|
+
def list_files(in_path):
|
|
13
|
+
img_files = []
|
|
14
|
+
mask_files = []
|
|
15
|
+
gt_files = []
|
|
16
|
+
for (dirpath, dirnames, filenames) in os.walk(in_path):
|
|
17
|
+
for file in filenames:
|
|
18
|
+
filename, ext = os.path.splitext(file)
|
|
19
|
+
ext = str.lower(ext)
|
|
20
|
+
if ext == '.jpg' or ext == '.jpeg' or ext == '.gif' or ext == '.png' or ext == '.pgm':
|
|
21
|
+
img_files.append(os.path.join(dirpath, file))
|
|
22
|
+
elif ext == '.bmp':
|
|
23
|
+
mask_files.append(os.path.join(dirpath, file))
|
|
24
|
+
elif ext == '.xml' or ext == '.gt' or ext == '.txt':
|
|
25
|
+
gt_files.append(os.path.join(dirpath, file))
|
|
26
|
+
elif ext == '.zip':
|
|
27
|
+
continue
|
|
28
|
+
# img_files.sort()
|
|
29
|
+
# mask_files.sort()
|
|
30
|
+
# gt_files.sort()
|
|
31
|
+
return img_files, mask_files, gt_files
|
|
32
|
+
|
|
33
|
+
def saveResult(img_file, img, boxes, dirname='./result/', verticals=None, texts=None):
|
|
34
|
+
""" save text detection result one by one
|
|
35
|
+
Args:
|
|
36
|
+
img_file (str): image file name
|
|
37
|
+
img (array): raw image context
|
|
38
|
+
boxes (array): array of result file
|
|
39
|
+
Shape: [num_detections, 4] for BB output / [num_detections, 4] for QUAD output
|
|
40
|
+
Return:
|
|
41
|
+
None
|
|
42
|
+
"""
|
|
43
|
+
img = np.array(img)
|
|
44
|
+
|
|
45
|
+
# make result file list
|
|
46
|
+
filename, file_ext = os.path.splitext(os.path.basename(img_file))
|
|
47
|
+
|
|
48
|
+
# result directory
|
|
49
|
+
res_file = dirname + "res_" + filename + '.txt'
|
|
50
|
+
res_img_file = dirname + "res_" + filename + '.jpg'
|
|
51
|
+
|
|
52
|
+
if not os.path.isdir(dirname):
|
|
53
|
+
os.mkdir(dirname)
|
|
54
|
+
|
|
55
|
+
with open(res_file, 'w') as f:
|
|
56
|
+
for i, box in enumerate(boxes):
|
|
57
|
+
poly = np.array(box).astype(np.int32).reshape((-1))
|
|
58
|
+
strResult = ','.join([str(p) for p in poly]) + '\r\n'
|
|
59
|
+
f.write(strResult)
|
|
60
|
+
|
|
61
|
+
poly = poly.reshape(-1, 2)
|
|
62
|
+
cv2.polylines(img, [poly.reshape((-1, 1, 2))], True, color=(0, 0, 255), thickness=2)
|
|
63
|
+
ptColor = (0, 255, 255)
|
|
64
|
+
if verticals is not None:
|
|
65
|
+
if verticals[i]:
|
|
66
|
+
ptColor = (255, 0, 0)
|
|
67
|
+
|
|
68
|
+
if texts is not None:
|
|
69
|
+
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
70
|
+
font_scale = 0.5
|
|
71
|
+
cv2.putText(img, "{}".format(texts[i]), (poly[0][0]+1, poly[0][1]+1), font, font_scale, (0, 0, 0), thickness=1)
|
|
72
|
+
cv2.putText(img, "{}".format(texts[i]), tuple(poly[0]), font, font_scale, (0, 255, 255), thickness=1)
|
|
73
|
+
|
|
74
|
+
# Save result image
|
|
75
|
+
cv2.imwrite(res_img_file, img)
|
|
76
|
+
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
3
|
+
MIT License
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
# -*- coding: utf-8 -*-
|
|
7
|
+
import numpy as np
|
|
8
|
+
from skimage import io
|
|
9
|
+
import cv2
|
|
10
|
+
|
|
11
|
+
def loadImage(img_file):
|
|
12
|
+
img = io.imread(img_file) # RGB order
|
|
13
|
+
if img.shape[0] == 2: img = img[0]
|
|
14
|
+
if len(img.shape) == 2 : img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
|
15
|
+
if img.shape[2] == 4: img = img[:,:,:3]
|
|
16
|
+
img = np.array(img)
|
|
17
|
+
|
|
18
|
+
return img
|
|
19
|
+
|
|
20
|
+
def normalizeMeanVariance(in_img, mean=(0.485, 0.456, 0.406), variance=(0.229, 0.224, 0.225)):
|
|
21
|
+
# should be RGB order
|
|
22
|
+
img = in_img.copy().astype(np.float32) / 255.0
|
|
23
|
+
|
|
24
|
+
img -= np.array(mean, dtype=np.float32)
|
|
25
|
+
img /= np.array(variance, dtype=np.float32)
|
|
26
|
+
return img
|
|
27
|
+
|
|
28
|
+
def denormalizeMeanVariance(in_img, mean=(0.485, 0.456, 0.406), variance=(0.229, 0.224, 0.225)):
|
|
29
|
+
# should be RGB order
|
|
30
|
+
img = in_img.copy()
|
|
31
|
+
img *= variance
|
|
32
|
+
img += mean
|
|
33
|
+
img *= 255.0
|
|
34
|
+
img = np.clip(img, 0, 255).astype(np.uint8)
|
|
35
|
+
return img
|
|
36
|
+
|
|
37
|
+
def resize_aspect_ratio(img, square_size=2560, interpolation=cv2.INTER_CUBIC, mag_ratio=1.0, target_size=None, round=True):
|
|
38
|
+
height, width = img.shape[0:2]
|
|
39
|
+
|
|
40
|
+
max_sz = max(height, width)
|
|
41
|
+
# magnify image size
|
|
42
|
+
if target_size is None:
|
|
43
|
+
target_size = mag_ratio * max_sz
|
|
44
|
+
|
|
45
|
+
# set original image size
|
|
46
|
+
if target_size > square_size:
|
|
47
|
+
target_size = square_size
|
|
48
|
+
|
|
49
|
+
ratio = target_size / max_sz
|
|
50
|
+
|
|
51
|
+
target_w = int(width * ratio)
|
|
52
|
+
|
|
53
|
+
# check rounded width and adjust ratio again
|
|
54
|
+
if round and target_w % 32 != 0:
|
|
55
|
+
# round width
|
|
56
|
+
target_w = target_w + (32 - target_w % 32)
|
|
57
|
+
ratio = target_w / width
|
|
58
|
+
|
|
59
|
+
target_h = int(height * ratio)
|
|
60
|
+
|
|
61
|
+
resized = cv2.resize(img, (target_w, target_h), interpolation=interpolation)
|
|
62
|
+
|
|
63
|
+
return resized, ratio
|
|
64
|
+
|
|
65
|
+
def cvt2HeatmapImg(img):
|
|
66
|
+
img = (np.clip(img, 0, 1) * 255).astype(np.uint8)
|
|
67
|
+
img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
|
|
68
|
+
return img
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
3
|
+
MIT License
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
# -*- coding: utf-8 -*-
|
|
7
|
+
import torch
|
|
8
|
+
import torch.nn as nn
|
|
9
|
+
import torch.nn.functional as F
|
|
10
|
+
from torch.autograd import Variable
|
|
11
|
+
from .basenet.vgg16_bn import init_weights
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class RefineNet(nn.Module):
|
|
15
|
+
def __init__(self):
|
|
16
|
+
super(RefineNet, self).__init__()
|
|
17
|
+
|
|
18
|
+
self.last_conv = nn.Sequential(
|
|
19
|
+
nn.Conv2d(34, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True),
|
|
20
|
+
nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True),
|
|
21
|
+
nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
self.aspp1 = nn.Sequential(
|
|
25
|
+
nn.Conv2d(64, 128, kernel_size=3, dilation=6, padding=6), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
26
|
+
nn.Conv2d(128, 128, kernel_size=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
27
|
+
nn.Conv2d(128, 1, kernel_size=1)
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
self.aspp2 = nn.Sequential(
|
|
31
|
+
nn.Conv2d(64, 128, kernel_size=3, dilation=12, padding=12), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
32
|
+
nn.Conv2d(128, 128, kernel_size=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
33
|
+
nn.Conv2d(128, 1, kernel_size=1)
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
self.aspp3 = nn.Sequential(
|
|
37
|
+
nn.Conv2d(64, 128, kernel_size=3, dilation=18, padding=18), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
38
|
+
nn.Conv2d(128, 128, kernel_size=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
39
|
+
nn.Conv2d(128, 1, kernel_size=1)
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
self.aspp4 = nn.Sequential(
|
|
43
|
+
nn.Conv2d(64, 128, kernel_size=3, dilation=24, padding=24), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
44
|
+
nn.Conv2d(128, 128, kernel_size=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
|
|
45
|
+
nn.Conv2d(128, 1, kernel_size=1)
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
init_weights(self.last_conv.modules())
|
|
49
|
+
init_weights(self.aspp1.modules())
|
|
50
|
+
init_weights(self.aspp2.modules())
|
|
51
|
+
init_weights(self.aspp3.modules())
|
|
52
|
+
init_weights(self.aspp4.modules())
|
|
53
|
+
|
|
54
|
+
def forward(self, y, upconv4):
|
|
55
|
+
refine = torch.cat([y.permute(0,3,1,2), upconv4], dim=1)
|
|
56
|
+
refine = self.last_conv(refine)
|
|
57
|
+
|
|
58
|
+
aspp1 = self.aspp1(refine)
|
|
59
|
+
aspp2 = self.aspp2(refine)
|
|
60
|
+
aspp3 = self.aspp3(refine)
|
|
61
|
+
aspp4 = self.aspp4(refine)
|
|
62
|
+
|
|
63
|
+
#out = torch.add([aspp1, aspp2, aspp3, aspp4], dim=1)
|
|
64
|
+
out = aspp1 + aspp2 + aspp3 + aspp4
|
|
65
|
+
return out.permute(0, 2, 3, 1) # , refine.permute(0,2,3,1)
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
__all__ = ["Logger"]
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class Logger(logging.Logger):
|
|
8
|
+
def __init__(self, name: str, level=None, fmt=None):
|
|
9
|
+
fmt = fmt or "CRAFT (%(levelname)s): %(message)s"
|
|
10
|
+
level = level or "INFO"
|
|
11
|
+
super().__init__(name, level)
|
|
12
|
+
handler = logging.StreamHandler()
|
|
13
|
+
formatter = logging.Formatter(fmt)
|
|
14
|
+
handler.setFormatter(formatter)
|
|
15
|
+
self.addHandler(handler)
|
|
16
|
+
|
|
17
|
+
def log_upload_success(self, name, target_path: str):
|
|
18
|
+
"""
|
|
19
|
+
Log (info) success info when the file(s) upload is done.
|
|
20
|
+
"""
|
|
21
|
+
self.info(f"Uploaded: `{name}` --> `{target_path}`")
|
|
@@ -0,0 +1,264 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
|
|
3
|
+
from shapely.geometry import Polygon
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
# from A1111 stable diffusion code with some modification.
|
|
7
|
+
class ReplaceHelper:
|
|
8
|
+
def __init__(self):
|
|
9
|
+
self.replaced = []
|
|
10
|
+
|
|
11
|
+
def replace(self, obj, field, func):
|
|
12
|
+
original = getattr(obj, field, None)
|
|
13
|
+
if original is None:
|
|
14
|
+
return None
|
|
15
|
+
|
|
16
|
+
self.replaced.append((obj, field, original))
|
|
17
|
+
setattr(obj, field, func)
|
|
18
|
+
|
|
19
|
+
return original
|
|
20
|
+
|
|
21
|
+
def restore(self):
|
|
22
|
+
for obj, field, original in self.replaced:
|
|
23
|
+
setattr(obj, field, original)
|
|
24
|
+
|
|
25
|
+
self.replaced.clear()
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class InitializeOnMeta(ReplaceHelper):
|
|
29
|
+
"""
|
|
30
|
+
Context manager that causes all parameters for linear/conv2d layers to be allocated on meta device,
|
|
31
|
+
which results in those parameters having no values and taking no memory. model.to() will be broken and
|
|
32
|
+
will need to be repaired by using LoadStateDictOnMeta below when loading params from state dict.
|
|
33
|
+
|
|
34
|
+
Usage:
|
|
35
|
+
```
|
|
36
|
+
with InitializeOnMeta():
|
|
37
|
+
...
|
|
38
|
+
```
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
def __enter__(self):
|
|
42
|
+
|
|
43
|
+
def set_device(x):
|
|
44
|
+
x["device"] = "meta"
|
|
45
|
+
return x
|
|
46
|
+
|
|
47
|
+
linear_init = self.replace(torch.nn.Linear, '__init__', lambda *args, **kwargs: linear_init(*args, **set_device(kwargs)))
|
|
48
|
+
conv2d_init = self.replace(torch.nn.Conv2d, '__init__', lambda *args, **kwargs: conv2d_init(*args, **set_device(kwargs)))
|
|
49
|
+
self.replace(torch.nn.Module, 'to', lambda *args, **kwargs: None)
|
|
50
|
+
|
|
51
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
52
|
+
self.restore()
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class LoadStateDictOnMeta(ReplaceHelper):
|
|
56
|
+
"""
|
|
57
|
+
Context manager that allows to read parameters from state_dict into a model that has some of its parameters in the meta device.
|
|
58
|
+
As those parameters are read from state_dict, they will be deleted from it, so by the end state_dict will be mostly empty, to save memory.
|
|
59
|
+
Meant to be used together with InitializeOnMeta above.
|
|
60
|
+
|
|
61
|
+
Usage:
|
|
62
|
+
```
|
|
63
|
+
with LoadStateDictOnMeta(state_dict):
|
|
64
|
+
model.load_state_dict(state_dict, strict=False)
|
|
65
|
+
```
|
|
66
|
+
"""
|
|
67
|
+
|
|
68
|
+
def __init__(self, state_dict, device, weight_dtype_conversion=None):
|
|
69
|
+
super().__init__()
|
|
70
|
+
self.state_dict = state_dict
|
|
71
|
+
self.device = device
|
|
72
|
+
self.weight_dtype_conversion = weight_dtype_conversion or {}
|
|
73
|
+
self.default_dtype = self.weight_dtype_conversion.get('', None)
|
|
74
|
+
|
|
75
|
+
def get_weight_dtype(self, key):
|
|
76
|
+
key_first_term, _ = key.split('.', 1)
|
|
77
|
+
return self.weight_dtype_conversion.get(key_first_term, self.default_dtype)
|
|
78
|
+
|
|
79
|
+
def __enter__(self):
|
|
80
|
+
|
|
81
|
+
sd = self.state_dict
|
|
82
|
+
device = self.device
|
|
83
|
+
|
|
84
|
+
def load_from_state_dict(original, module, state_dict, prefix, *args, **kwargs):
|
|
85
|
+
used_param_keys = []
|
|
86
|
+
|
|
87
|
+
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d, torch.nn.GroupNorm,)):
|
|
88
|
+
# HACK add assign=True to local_metadata for some cases
|
|
89
|
+
args[0]['assign_to_params_buffers'] = True
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
for name, param in module._parameters.items():
|
|
93
|
+
if param is None:
|
|
94
|
+
continue
|
|
95
|
+
|
|
96
|
+
key = prefix + name
|
|
97
|
+
sd_param = sd.pop(key, None)
|
|
98
|
+
if sd_param is not None:
|
|
99
|
+
dtype = self.get_weight_dtype(key)
|
|
100
|
+
if dtype is None:
|
|
101
|
+
state_dict[key] = sd_param
|
|
102
|
+
else:
|
|
103
|
+
state_dict[key] = sd_param.to(dtype=dtype)
|
|
104
|
+
used_param_keys.append(key)
|
|
105
|
+
|
|
106
|
+
if param.is_meta:
|
|
107
|
+
dtype = sd_param.dtype if sd_param is not None else param.dtype
|
|
108
|
+
module._parameters[name] = torch.nn.parameter.Parameter(torch.empty_like(param, device=device, dtype=dtype), requires_grad=param.requires_grad)
|
|
109
|
+
|
|
110
|
+
for name in module._buffers:
|
|
111
|
+
key = prefix + name
|
|
112
|
+
|
|
113
|
+
sd_param = sd.pop(key, None)
|
|
114
|
+
if sd_param is not None:
|
|
115
|
+
state_dict[key] = sd_param
|
|
116
|
+
used_param_keys.append(key)
|
|
117
|
+
|
|
118
|
+
original(module, state_dict, prefix, *args, **kwargs)
|
|
119
|
+
|
|
120
|
+
for key in used_param_keys:
|
|
121
|
+
state_dict.pop(key, None)
|
|
122
|
+
|
|
123
|
+
def load_state_dict(original, module, state_dict, strict=True):
|
|
124
|
+
"""torch makes a lot of copies of the dictionary with weights, so just deleting entries from state_dict does not help
|
|
125
|
+
because the same values are stored in multiple copies of the dict. The trick used here is to give torch a dict with
|
|
126
|
+
all weights on meta device, i.e. deleted, and then it doesn't matter how many copies torch makes.
|
|
127
|
+
|
|
128
|
+
In _load_from_state_dict, the correct weight will be obtained from a single dict with the right weights (sd).
|
|
129
|
+
|
|
130
|
+
The dangerous thing about this is if _load_from_state_dict is not called, (if some exotic module overloads
|
|
131
|
+
the function and does not call the original) the state dict will just fail to load because weights
|
|
132
|
+
would be on the meta device.
|
|
133
|
+
"""
|
|
134
|
+
|
|
135
|
+
if state_dict is sd:
|
|
136
|
+
state_dict = {k: v.to(device="meta", dtype=v.dtype) for k, v in state_dict.items()}
|
|
137
|
+
|
|
138
|
+
original(module, state_dict, strict=strict)
|
|
139
|
+
|
|
140
|
+
module_load_state_dict = self.replace(torch.nn.Module, 'load_state_dict', lambda *args, **kwargs: load_state_dict(module_load_state_dict, *args, **kwargs))
|
|
141
|
+
module_load_from_state_dict = self.replace(torch.nn.Module, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(module_load_from_state_dict, *args, **kwargs))
|
|
142
|
+
linear_load_from_state_dict = self.replace(torch.nn.Linear, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(linear_load_from_state_dict, *args, **kwargs))
|
|
143
|
+
conv2d_load_from_state_dict = self.replace(torch.nn.Conv2d, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(conv2d_load_from_state_dict, *args, **kwargs))
|
|
144
|
+
layer_norm_load_from_state_dict = self.replace(torch.nn.LayerNorm, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(layer_norm_load_from_state_dict, *args, **kwargs))
|
|
145
|
+
group_norm_load_from_state_dict = self.replace(torch.nn.GroupNorm, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(group_norm_load_from_state_dict, *args, **kwargs))
|
|
146
|
+
|
|
147
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
148
|
+
self.restore()
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def get_center(box):
|
|
152
|
+
"""사각형 박스(box)의 중심 좌표 (cx, cy) 반환"""
|
|
153
|
+
x_coords = [point[0] for point in box]
|
|
154
|
+
y_coords = [point[1] for point in box]
|
|
155
|
+
return (sum(x_coords) / 4, sum(y_coords) / 4)
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def get_box_center_y(box):
|
|
159
|
+
"""각 박스의 중심 y 좌표 반환"""
|
|
160
|
+
return box[:, 1].mean()
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def merge_boxes(boxA, boxB):
|
|
164
|
+
# 두 박스의 모든 점을 하나의 배열로 합치기
|
|
165
|
+
points = np.vstack([boxA, boxB])
|
|
166
|
+
|
|
167
|
+
print("<< from box ", boxA, boxB)
|
|
168
|
+
|
|
169
|
+
# Convex Hull 적용
|
|
170
|
+
hull = cv2.convexHull(points.astype(np.float32))
|
|
171
|
+
|
|
172
|
+
# 최소 면적 회전된 사각형 찾기
|
|
173
|
+
rect = cv2.minAreaRect(hull) # (center (x, y), (w, h), angle)
|
|
174
|
+
|
|
175
|
+
# 회전된 사각형을 4개의 꼭짓점으로 변환
|
|
176
|
+
box = cv2.boxPoints(rect)
|
|
177
|
+
box = np.int32(box) # 정수 좌표 변환
|
|
178
|
+
|
|
179
|
+
# make clock-wise order
|
|
180
|
+
startidx = box.sum(axis=1).argmin()
|
|
181
|
+
box = np.roll(box, 4-startidx, 0)
|
|
182
|
+
box = np.array(box)
|
|
183
|
+
print(">> to box -> ", box)
|
|
184
|
+
|
|
185
|
+
return box
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def sort_text_boxes(boxes, y_threshold=5, merge=False):
|
|
189
|
+
"""
|
|
190
|
+
CRAFT 결과로 얻은 텍스트 박스를 줄 단위로 정렬
|
|
191
|
+
:param boxes: 네 개의 좌표 (x, y)로 구성된 텍스트 박스 리스트
|
|
192
|
+
:param y_threshold: 같은 줄로 묶을 y 축 거리 기준
|
|
193
|
+
:return: 정렬된 텍스트 박스 리스트 (줄 단위)
|
|
194
|
+
"""
|
|
195
|
+
# y 값을 기준으로 정렬
|
|
196
|
+
sorted_boxes = sorted(boxes, key=get_box_center_y)
|
|
197
|
+
|
|
198
|
+
min_area = y_threshold * y_threshold * 1.2 # minimal area
|
|
199
|
+
|
|
200
|
+
# 줄 단위로 그룹화
|
|
201
|
+
lines = []
|
|
202
|
+
current_line = [sorted_boxes[0]]
|
|
203
|
+
|
|
204
|
+
for i in range(1, len(sorted_boxes)):
|
|
205
|
+
box = sorted_boxes[i]
|
|
206
|
+
curr_y = np.mean([get_box_center_y(b) for b in current_line])
|
|
207
|
+
box_y = get_box_center_y(box)
|
|
208
|
+
|
|
209
|
+
if len(current_line) > 0 and abs(box_y - curr_y) < y_threshold:
|
|
210
|
+
current_line.append(box)
|
|
211
|
+
else:
|
|
212
|
+
lines.append(current_line)
|
|
213
|
+
current_line = [box]
|
|
214
|
+
|
|
215
|
+
if current_line:
|
|
216
|
+
lines.append(current_line)
|
|
217
|
+
|
|
218
|
+
# 각 줄을 x 좌표 기준으로 정렬
|
|
219
|
+
sorted_lines = [sorted(line, key=lambda b: b[0][0]) for line in lines]
|
|
220
|
+
|
|
221
|
+
# 겹치거나 가까운 박스 검사
|
|
222
|
+
if not merge:
|
|
223
|
+
return sorted_lines
|
|
224
|
+
|
|
225
|
+
for j, lines in enumerate(sorted_lines):
|
|
226
|
+
current_line = [lines[0]]
|
|
227
|
+
for i in range(1, len(lines)):
|
|
228
|
+
last_box = current_line[-1]
|
|
229
|
+
box = lines[i]
|
|
230
|
+
|
|
231
|
+
last_poly = Polygon(last_box)
|
|
232
|
+
poly = Polygon(box)
|
|
233
|
+
|
|
234
|
+
intersect = poly.intersection(last_poly)
|
|
235
|
+
if not intersect.is_empty or poly.area < min_area or last_poly.area < min_area:
|
|
236
|
+
box = merge_boxes(last_box, box)
|
|
237
|
+
current_line[-1] = box
|
|
238
|
+
continue
|
|
239
|
+
|
|
240
|
+
if box[0][0] > last_box[1][0] and box[0][0] - last_box[1][0] < y_threshold * 0.3:
|
|
241
|
+
box = merge_boxes(last_box, box)
|
|
242
|
+
current_line[-1] = box
|
|
243
|
+
print("XXX mergable", i - 1, i)
|
|
244
|
+
continue
|
|
245
|
+
|
|
246
|
+
current_line.append(box)
|
|
247
|
+
|
|
248
|
+
sorted_lines[j] = current_line
|
|
249
|
+
|
|
250
|
+
return sorted_lines
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
def _get_box_height(box):
|
|
254
|
+
"""주어진 4개 좌표의 텍스트 박스에서 높이를 계산"""
|
|
255
|
+
y_coords = [point[1] for point in box]
|
|
256
|
+
height = (max(y_coords) - min(y_coords)) # 높이 = 최대 y값 - 최소 y값
|
|
257
|
+
return height
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
def get_average_height(boxes):
|
|
261
|
+
"""Calculate the average height of all text boxes"""
|
|
262
|
+
heights = [_get_box_height(box) for box in boxes]
|
|
263
|
+
avg_height = np.mean(heights) if heights else 0
|
|
264
|
+
return avg_height
|
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: craft-detector
|
|
3
|
+
Version: 0.2.1
|
|
4
|
+
Summary: Python package for Official implementation of Character Region Awareness for Text Detection (CRAFT)
|
|
5
|
+
Author-email: "Clova AI Research, NAVER Corp." <youngmin.baek@navercorp.com>
|
|
6
|
+
Maintainer-email: Won-Kyu Park <wkpark@gmail.com>
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Keywords: craft,deep-learning,machine-learning,ml,pytorch,text,text-detection
|
|
10
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Topic :: Education
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering
|
|
17
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Image Recognition
|
|
19
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
20
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
21
|
+
Requires-Python: >=3.10
|
|
22
|
+
Requires-Dist: opencv-python>=4.10.0
|
|
23
|
+
Requires-Dist: scikit-image>=0.25.1
|
|
24
|
+
Requires-Dist: scipy>=1.14.0
|
|
25
|
+
Requires-Dist: torch>=2.1.2
|
|
26
|
+
Requires-Dist: torchvision>=0.17.2
|
|
27
|
+
Description-Content-Type: text/markdown
|
|
28
|
+
|
|
29
|
+
## CRAFT: Character-Region Awareness For Text detection
|
|
30
|
+
Official Pytorch implementation of CRAFT text detector | [Paper](https://arxiv.org/abs/1904.01941) | [Pretrained Model](https://drive.google.com/open?id=1Jk4eGD7crsqCCg9C9VjCLkMN3ze8kutZ) | [Supplementary](https://youtu.be/HI8MzpY8KMI)
|
|
31
|
+
|
|
32
|
+
**[Youngmin Baek](mailto:youngmin.baek@navercorp.com), Bado Lee, Dongyoon Han, Sangdoo Yun, Hwalsuk Lee.**
|
|
33
|
+
|
|
34
|
+
Clova AI Research, NAVER Corp.
|
|
35
|
+
|
|
36
|
+
### Sample Results
|
|
37
|
+
|
|
38
|
+
### Overview
|
|
39
|
+
PyTorch implementation for CRAFT text detector that effectively detect text area by exploring each character region and affinity between characters. The bounding box of texts are obtained by simply finding minimum bounding rectangles on binary map after thresholding character region and affinity scores.
|
|
40
|
+
|
|
41
|
+
<img width="1000" alt="teaser" src="./figures/craft_example.gif">
|
|
42
|
+
|
|
43
|
+
## Updates
|
|
44
|
+
**13 Jun, 2019**: Initial update
|
|
45
|
+
**20 Jul, 2019**: Added post-processing for polygon result
|
|
46
|
+
**28 Sep, 2019**: Added the trained model on IC15 and the link refiner
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
## Getting started
|
|
50
|
+
### Install dependencies
|
|
51
|
+
#### Requirements
|
|
52
|
+
- PyTorch>=2.1.0
|
|
53
|
+
- torchvision>=0.17.2
|
|
54
|
+
- opencv-python>=4.10.0
|
|
55
|
+
- check requiremtns.txt
|
|
56
|
+
```
|
|
57
|
+
pip install -r requirements.txt
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
### Training
|
|
61
|
+
The code for training is not included in this repository, and we cannot release the full training code for IP reason.
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
### Test instruction using pretrained model
|
|
65
|
+
- Download the trained models
|
|
66
|
+
|
|
67
|
+
*Model name* | *Used datasets* | *Languages* | *Purpose* | *Model Link* |
|
|
68
|
+
| :--- | :--- | :--- | :--- | :--- |
|
|
69
|
+
General | SynthText, IC13, IC17 | Eng + MLT | For general purpose | [Click](https://drive.google.com/open?id=1Jk4eGD7crsqCCg9C9VjCLkMN3ze8kutZ)
|
|
70
|
+
IC15 | SynthText, IC15 | Eng | For IC15 only | [Click](https://drive.google.com/open?id=1i2R7UIUqmkUtF0jv_3MXTqmQ_9wuAnLf)
|
|
71
|
+
LinkRefiner | CTW1500 | - | Used with the General Model | [Click](https://drive.google.com/open?id=1XSaFwBkOaFOdtk4Ane3DFyJGPRw6v5bO)
|
|
72
|
+
|
|
73
|
+
* Run with pretrained model
|
|
74
|
+
``` (with python 3.7)
|
|
75
|
+
python test.py --trained_model=[weightfile] --test_folder=[folder path to test images]
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
The result image and socre maps will be saved to `./result` by default.
|
|
79
|
+
|
|
80
|
+
### Arguments
|
|
81
|
+
* `--trained_model`: pretrained model
|
|
82
|
+
* `--text_threshold`: text confidence threshold
|
|
83
|
+
* `--low_text`: text low-bound score
|
|
84
|
+
* `--link_threshold`: link confidence threshold
|
|
85
|
+
* `--gpu`: use GPU for inference (default:False)
|
|
86
|
+
* `--canvas_size`: max image size for inference
|
|
87
|
+
* `--mag_ratio`: image magnification ratio
|
|
88
|
+
* `--poly`: enable polygon type result
|
|
89
|
+
* `--show_time`: show processing time
|
|
90
|
+
* `--test_folder`: folder path to input images
|
|
91
|
+
* `--refine`: use link refiner for sentense-level dataset
|
|
92
|
+
* `--refiner_model`: pretrained refiner model
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
## Links
|
|
96
|
+
- Repo of recognition : https://github.com/clovaai/deep-text-recognition-benchmark
|
|
97
|
+
|
|
98
|
+
## Citation
|
|
99
|
+
```
|
|
100
|
+
@inproceedings{baek2019character,
|
|
101
|
+
title={Character Region Awareness for Text Detection},
|
|
102
|
+
author={Baek, Youngmin and Lee, Bado and Han, Dongyoon and Yun, Sangdoo and Lee, Hwalsuk},
|
|
103
|
+
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
|
|
104
|
+
pages={9365--9374},
|
|
105
|
+
year={2019}
|
|
106
|
+
}
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
## License
|
|
110
|
+
```
|
|
111
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
112
|
+
|
|
113
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
114
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
115
|
+
in the Software without restriction, including without limitation the rights
|
|
116
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
117
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
118
|
+
furnished to do so, subject to the following conditions:
|
|
119
|
+
|
|
120
|
+
The above copyright notice and this permission notice shall be included in
|
|
121
|
+
all copies or substantial portions of the Software.
|
|
122
|
+
|
|
123
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
124
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
125
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
126
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
127
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
128
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
129
|
+
THE SOFTWARE.
|
|
130
|
+
```
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
craft_detector/__init__.py,sha256=J2JWGXZLeBYPRXIqHh95jDBmN_NDjkHxJWWAE1wUQJg,44
|
|
2
|
+
craft_detector/craft.py,sha256=G0FAiO0cM8sa5OVH-F6sjJ_-RfBUzRrAz2Lf3QziqpQ,2663
|
|
3
|
+
craft_detector/craft_utils.py,sha256=QCHfzkswREdxesS4f3lMJBb4VCy514e8oKx9fFTSpd4,11715
|
|
4
|
+
craft_detector/file_utils.py,sha256=3aCrxQSzVqN-fYb6-ortBfLtS00mbvncAP9JVpppT5w,2877
|
|
5
|
+
craft_detector/imgproc.py,sha256=OaldhhfXncwVvaYon5x3Una9JNLiUGG63fe0DS-8pII,1903
|
|
6
|
+
craft_detector/refinenet.py,sha256=KvM_hBcKCPMOiLAQSWOzhh7EmfzJWZkOisQQZk5heOk,2553
|
|
7
|
+
craft_detector/basenet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
8
|
+
craft_detector/basenet/vgg16_bn.py,sha256=UQAn2f9AZtZ4R2lGeUnE8hsXhvniixwe-UwqSik_SOU,2857
|
|
9
|
+
craft_detector/utils/logging.py,sha256=7-ZDv1hJQDlIUaAT2-Cl8F2PsX-f363MUs4AQ-V_ttg,622
|
|
10
|
+
craft_detector/utils/utils.py,sha256=GUwVXFeQH-7ZYAwDDRyhPMd2hMIGqNc0LCuOnQv7UDA,10015
|
|
11
|
+
craft_detector-0.2.1.dist-info/METADATA,sha256=PzoH_Tz78SixipkRMAHP6dHvYrCHg0HKZa_MT8SRFSE,5445
|
|
12
|
+
craft_detector-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
13
|
+
craft_detector-0.2.1.dist-info/licenses/LICENSE,sha256=v-zJfkiIiWt6zwhTQ3K68leWuYwIe6WkJLRu4lreyHE,1064
|
|
14
|
+
craft_detector-0.2.1.dist-info/RECORD,,
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
Copyright (c) 2019-present NAVER Corp.
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
|
|
10
|
+
The above copyright notice and this permission notice shall be included in
|
|
11
|
+
all copies or substantial portions of the Software.
|
|
12
|
+
|
|
13
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
14
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
15
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
16
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
17
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
18
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
19
|
+
THE SOFTWARE.
|