一、開發背景
圖像中的文本識別近幾年來備受矚目。通常來說,圖片中的文本能夠比圖片中其他內容提供更加豐富的信息。因此,圖像文本識別能夠將圖像中的文本區域轉化成計算機可以讀取和編輯的符號,打通了從圖像到文本再到信息的通路。
隨着計算機算力的提升,基於深度學習方法的本文識別技術逐漸成為主流,而深度學習中數據集的獲取是重中之重。本腳本實現讀取語料集中的文本內容,以保存為圖像形式的數據集,用於模型訓練。
二、腳本效果
1、IDE中的運行界面
(1)選擇字體文件
(2)生成數據集
2、生成的圖像
不使用數據增強
使用數據增強
3、映射表
存儲圖像文件名和類別序列標注的對應關系
三、具體開發
1、功能需求
- 根據用戶指定的語料數據生成圖像文件及映射表
- 用戶可自行更改文本長度,圖像數量及圖像尺寸
- 用戶可自行選擇是否進行增強處理
2、實際項目
1. 項目結構
(1)根目錄下的fonts文件夾用於存放ttf字體文件, imageset文件夾用於存放輸出圖像和映射表
(2)config中設置相關參數並存放語料文件, dict5990.txt是字典, sentences.txt是語料集
2. 實現思路
3. 代碼實現
1. 設置參數
# 語料集
corpus = 'config/sentences.txt'
dict = 'config/dict5990.txt'
# 字體文件路徑
FONT_PATH = 'fonts/'
# 輸出路徑
OUTPUT_DIR = 'imageset/'
# 樣本總數
n_samples = 50
# 每行最大長度
sentence_lim = 10
# 畫布能容納的最大序列長度,對應img_w
canvas_lim = 50
2. 構建生成器
1. 加載字體文件
# 選擇字體
root = tk.Tk()
root.withdraw()
self.font_path = filedialog.askopenfilename()
def load_fonts(self, factor, font_path):
""" 加載字體文件並設定字體大小
"""
self.fonts = []
# 加載字體文件
font = ImageFont.truetype(font_path, int(self.img_h*factor), 0)
self.fonts.append(font)
2. 構建字典
def build_dict(self):
""" 打開字典,加載全部字符到list
每行是一個字
"""
with codecs.open(self.dictfile, mode='r', encoding='utf-8') as f:
# 按行讀取語料
for line in f:
# 當前行單詞去除結尾,為了正常讀取空格,第一行兩個空格
word = line.strip('\r\n')
# 只要沒超出上限就繼續添加單詞
self.dict.append(word)
# 最后一位作為空白符
self.blank_label = len(self.dict)
3. 加載語料
def build_train_list(self, num_rows, max_row_len=None):
# 過濾語料,留下適合的內容組成訓練list
assert max_row_len <= self.img_lim
self.num_rows = num_rows
self.max_row_len = max_row_len
sentence_list = []
self.train_list = []
with codecs.open(self.corpus_file, mode='r', encoding='utf-8') as f:
# 按行讀取語料
for line in f:
sentence = line.rstrip().replace(' ', '') # 當前行單詞
if len(sentence) <= max_row_len and len(sentence_list) < num_rows:
# 只要句子長度不超過畫布上限且句子數量沒超出上限就繼續添加
sentence_list.append(sentence)
elif len(sentence) > max_row_len and len(sentence_list) < num_rows:
# 截斷句子
sentence_list.append(sentence[0:max_row_len])
if len(sentence_list) < self.num_rows:
raise IOError('語料不夠')
for i, sentence in enumerate(sentence_list):
# 遍歷語料中的每一句(行)
# 將單詞分成字符,然后找到每個字符對應的整數ID list
label_sequence = []
for j, word in enumerate(sentence): # 檢查句子中是否包含生僻字
try:
index = self.dict.index(word)
label_sequence.append(index)
except ValueError:
print("字典不包含:{},已忽略".format(word))
sentence_list[i] = sentence_list[i][0:j] + sentence_list[i][j+1:] # 從該句中刪除生僻字
self.train_list = sentence_list # 過濾后的訓練集
np.random.shuffle(self.train_list) # 打亂順序
4. 保存映射表
def mapping_list(self):
# 寫圖像文件名和類別序列的對照表
file_path = os.path.join(cfg.OUTPUT_DIR, 'map_list.txt')
with codecs.open(file_path, mode='w', encoding='utf-8') as f:
for i in range(len(self.train_list)):
f.write("{}.png {} \n".format(i, self.train_list[i]))
5. 繪制圖像
def paint_text(self, text, i):
""" 使用PIL繪制文本圖像,傳入畫布尺寸,返回文本圖像
:param h: 畫布高度
:param w: 畫布寬度
"""
# 創建畫布
canvas = np.zeros(shape=(self.img_h, self.img_w), dtype=np.uint8)
canvas[0:] = 255
# 轉換圖像模式,保證合成的兩張圖尺寸模式一致
ndimg = Image.fromarray(canvas).convert('RGBA')
draw = ImageDraw.Draw(ndimg)
font = self.fonts[-1]
text_size = font.getsize(text) # 獲取當前字體下的文本區域大小
# 自動調整字體大小避免超出邊界, 至少留白水平20%
margin = [self.img_w - int(0.2*self.img_w), self.img_h - int(0.2*self.img_h)]
while (text_size[0] > margin[0]) or (text_size[1] > margin[1]):
self.font_factor -= 0.1
self.load_fonts(self.font_factor, self.font_path)
font = self.fonts[-1]
text_size = font.getsize(text)
# 隨機平移
horizontal_space = self.img_w - text_size[0]
vertical_space = self.img_h - text_size[1]
start_x = np.random.randint(2, horizontal_space-2)
start_y = np.random.randint(2, vertical_space-2)
# 繪制當前文本行
draw.text((start_x, start_y), text, font=font, fill=(0, 0, 0, 255))
img_array = np.array(ndimg)
# 轉灰度圖
grey_img = img_array[:, :, 0] # [32, 256, 4]
if self.aug == True:
auged = augmentation(grey_img)
ndimg = Image.fromarray(auged).convert('RGBA')
save_path = os.path.join(cfg.OUTPUT_DIR, '{}.png'.format(i)) # 類別序列即文件名
ndimg.save(save_path)
6. 數據增強
def speckle(img):
severity = np.random.uniform(0, 0.6*255)
blur = ndimage.gaussian_filter(np.random.randn(*img.shape) * severity, 1)
img_speck = (img + blur)
img_speck[img_speck > 255] = 255
img_speck[img_speck <= 0] = 0
return img_speck
def augmentation(img, ):
# 不能直接在原始image上改動
image = img.copy()
img_h, img_w = img.shape
mode = np.random.randint(0, 9)
'''添加隨機模糊和噪聲'''
# 高斯模糊
if mode == 0:
image = cv2.GaussianBlur(image,(5, 5), np.random.randint(1, 10))
# 模糊后二值化,虛化邊緣
if mode == 1:
image = cv2.GaussianBlur(image, (9, 9), np.random.randint(1, 8))
ret, th = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
thresh = image.copy()
thresh[thresh >= th] = 0
thresh[thresh < th] = 255
image = thresh
# 橫線干擾
if mode == 2:
for i in range(0, img_w, 2):
cv2.line(image, (0, i), (img_w, i), 0, 1)
# 豎線
if mode == 3:
for i in range(0, img_w, 2):
cv2.line(image, (i, 0), (i, img_h), 0, 1)
# 十字線
if mode == 4:
for i in range(0, img_h, 2):
cv2.line(image, (0, i), (img_w, i), 0, 1)
for i in range(0, img_w, 2):
cv2.line(image, (i, 0), (i, img_h), 0, 1)
# 左右運動模糊
if mode == 5:
kernel_size = 5
kernel_motion_blur = np.zeros((kernel_size, kernel_size))
kernel_motion_blur[int((kernel_size - 1) / 2), :] = np.ones(kernel_size)
kernel_motion_blur = kernel_motion_blur / kernel_size
image = cv2.filter2D(image, -1, kernel_motion_blur)
# 上下運動模糊
if mode == 6:
kernel_size = 9
kernel_motion_blur = np.zeros((kernel_size, kernel_size))
kernel_motion_blur[:, int((kernel_size - 1) / 2)] = np.ones(kernel_size)
kernel_motion_blur = kernel_motion_blur / kernel_size
image = cv2.filter2D(image, -1, kernel_motion_blur)
# 高斯噪聲
if mode == 7:
row, col = [img_h, img_w]
mean = 0
sigma = 1
gauss = np.random.normal(mean, sigma, (row, col))
gauss = gauss.reshape(row, col)
noisy = image + gauss
image = noisy.astype(np.uint8)
# 污跡
if mode == 8:
image = speckle(image)
return image
4. 使用說明
運行sample_generator.py后會跳出對話框, 選擇字體文件即可生成數據集
從文本到圖像——文本識別數據集生成器
注:本文著作權歸作者,由demo大師代發,拒絕轉載,轉載需要作者授權