forked from NVlabs/ocropus3-ocroseg
-
Notifications
You must be signed in to change notification settings - Fork 0
/
degrade-pages
executable file
·128 lines (106 loc) · 3.78 KB
/
degrade-pages
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/python
import os
import os.path
import argparse
import simplejson
import scipy.ndimage as ndi
from pylab import *
from dlinputs import gopen, utils, filters
from ocroseg import degrade
default_degrade = "translation=0.03, rotation=1.0, scale=0.03, aniso=0.03"
parser = argparse.ArgumentParser(
"Degrade pages for page segmentation training.")
parser.add_argument("-d", "--degrade", default=default_degrade)
parser.add_argument("-f", "--fields", default="framed.png,lines.png")
parser.add_argument("--maxdelta", default="0.0,5.0")
parser.add_argument("--distort", default="20.0:100.0")
parser.add_argument("--noisesigma", default="0.5:1.5")
parser.add_argument("--noisemag", default="0.1:0.2")
parser.add_argument("--blur", default="0.0,1.0")
parser.add_argument("--blobsper100", default="0.0,0.3")
parser.add_argument("--blobsize", default="1.0:20.0")
parser.add_argument("--rotation_limit", type=float, default=0.1)
parser.add_argument("--display", type=int, default=0)
parser.add_argument("--limit", type=int, default=1000000000)
parser.add_argument("input")
parser.add_argument("output")
args = parser.parse_args()
if args.display > 0:
ion()
ifield, ofield = args.fields.split(",")
assert not os.path.exists(args.output)
degradation = eval("dict(%s)" % args.degrade)
def R(r):
if "," in r:
lo, hi = [float(x) for x in r.split(",")]
return np.random.uniform(lo, hi)
elif ":" in r:
lo, hi = [float(x) for x in r.split(":")]
return np.exp(np.random.uniform(np.log(lo), np.log(hi)))
else:
return float(r)
def degrade_sample(sample):
input = sample["input"]
output = sample["output"]
h, w = input.shape
#print "*blobs"
nblobs = int(R(args.blobsper100) * h * w / (100.0 * 100.0))
input = np.maximum(input, degrade.random_blobs(
(h, w), nblobs, R(args.blobsize)))
#print "*gauss"
noisesig = R(args.noisesigma)
noisemag = R(args.noisemag)
input = degrade.gauss_degrade(input, noisesig, noisemag)
#print "*distort"
maxdelta = R(args.maxdelta)
distort = R(args.distort)
input, output = degrade.random_distort([input, output],
maxdelta=maxdelta,
sigma=distort)
#print "*trs"
f, params = degrade.random_trs(rotation=args.rotation_limit)
input, output = f(input), f(output)
input = clip(input, 0, 1)
output = clip(output, 0, 1)
#print "*blur"
blur = R(args.blur)
if blur > 1e-5:
input = ndi.gaussian_filter(input, blur)
if amax(input) > 0:
input /= amax(input)
#print "*done"
params.update(dict(blur=blur,
nblobs=nblobs,
maxdelta=maxdelta,
distort=distort,
noisesig=noisesig,
noisemag=noisemag))
return dict(input=input, output=output, params=params)
pipeline = filters.compose(
filters.ren({"input": ifield, "output": ofield}),
filters.transform(degrade_sample),
filters.ren({ifield: "input", ofield: "output", "params.json": "params"}),
)
source = gopen.sharditerator_once(args.input)
source = pipeline(source)
sink = gopen.open_sink(args.output)
for i, sample in enumerate(source):
if i >= args.limit:
break
if i % 1 == 0:
print i, sample.get("__key__")
if i == 0:
utils.print_sample(sample)
if args.display > 0 and i % args.display == 0:
utils.print_sample(sample)
print simplejson.dumps(sample["params.json"], indent=4)
clf()
subplot(121)
imshow(sample[ifield], cmap=cm.gray, vmin=0, vmax=1)
subplot(122)
imshow(sample[ofield], cmap=cm.gray, vmin=0, vmax=1)
show()
ginput(1, 0.001)
print
sink.write(sample)
sink.close()