-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
79 lines (64 loc) · 2.51 KB
/
utils.py
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
"""
Utility functions
"""
import os
import numpy as np
from PIL import Image
def load_dict(lang):
"""Read the dictionnary file and returns all words in it.
"""
lang_dict = []
with open(
os.path.join(os.path.dirname(__file__), "dicts", lang + ".txt"),
"r",
encoding="utf8",
errors="ignore",
) as d:
lang_dict = [l for l in d.read().splitlines() if len(l) > 0]
return lang_dict
def load_fonts(lang):
"""Load all fonts in the fonts directories
"""
if lang in os.listdir(os.path.join(os.path.dirname(__file__), "fonts")):
return [
os.path.join(os.path.dirname(__file__), "fonts/{}".format(lang), font)
for font in os.listdir(
os.path.join(os.path.dirname(__file__), "fonts/{}".format(lang))
)
]
else:
return [
os.path.join(os.path.dirname(__file__), "fonts/latin", font)
for font in os.listdir(
os.path.join(os.path.dirname(__file__), "fonts/latin")
)
]
def salt_and_pepper_noise(image, prob=0.05):
# If the specified `prob` is negative or zero, we don't need to do anything.
if prob <= 0:
return image
arr = np.asarray(image)
original_dtype = arr.dtype
# Derive the number of intensity levels from the array datatype.
intensity_levels = 2 ** (arr[0, 0].nbytes * 8)
min_intensity = 0
max_intensity = intensity_levels - 1
# Generate an array with the same shape as the image's:
# Each entry will have:
# 1 with probability: 1 - prob
# 0 or np.nan (50% each) with probability: prob
random_image_arr = np.random.choice(
[min_intensity, 1, np.nan], p=[prob / 2, 1 - prob, prob / 2], size=arr.shape
)
# This results in an image array with the following properties:
# - With probability 1 - prob: the pixel KEEPS ITS VALUE (it was multiplied by 1)
# - With probability prob/2: the pixel has value zero (it was multiplied by 0)
# - With probability prob/2: the pixel has value np.nan (it was multiplied by np.nan)
# We need to to `arr.astype(np.float)` to make sure np.nan is a valid value.
salt_and_peppered_arr = arr.astype(np.float) * random_image_arr
# Since we want SALT instead of NaN, we replace it.
# We cast the array back to its original dtype so we can pass it to PIL.
salt_and_peppered_arr = np.nan_to_num(
salt_and_peppered_arr, nan=max_intensity
).astype(original_dtype)
return Image.fromarray(salt_and_peppered_arr)