EAST模型文本检测和识别代码解析

该代码包含使用 EAST 模型预测图像中文本和使用 CTC 损失训练单独文本识别模型的函数。

predict_img 函数接收图像路径、预训练的 EAST 模型路径和可选的像素阈值。它调整图像大小并应用预处理，然后使用 EAST 模型预测图像中文本区域的位置和方向。然后，它在预测的文本区域周围绘制边界框并将带注释的图像保存到以“_act.jpg”为后缀的文件中。它还保存一个预测的四边形区域以红色突出显示的图像版本，以“_predict.jpg”为后缀。如果在预测的区域内成功识别出任何文本，它将相应的坐标保存到与原始图像同名但以“.txt”为后缀的文本文件中。

predict_txt 函数类似于 predict_img，但它不绘制边界框，而是返回图像中预测的四边形区域列表。

其余代码用于使用 CTC 损失函数训练文本识别模型。它定义了一些常量和用于编码和解码文本标签的字典，以及输入图像的大小、批次大小和训练的时期数。train 函数未显示，但它可能将标记的图像数据集作为输入并训练一个模型来识别这些图像中的文本。

import numpy as np
from PIL import Image, ImageDraw
from keras.preprocessing import image
from keras.applications.vgg16 import preprocess_input

from east import cfg
from east.nms import nms
from east.eastnet import EAST
from east.utils import resize_image, sigmoid


east = EAST()
east_detect = east.east_network()


def predict_img(img_path, model_path, pixel_threshold=cfg.pixel_threshold, quiet=False):
    east_detect.load_weights(model_path)
    img = image.load_img(img_path)
    d_wight, d_height = resize_image(img, cfg.max_predict_img_size)
    img = img.resize((d_wight, d_height), Image.NEAREST).convert('RGB')
    img = image.img_to_array(img)
    img = preprocess_input(img, mode='tf')
    x = np.expand_dims(img, axis=0)
    y = east_detect.predict(x)

    y = np.squeeze(y, axis=0)
    y[:, :, :3] = sigmoid(y[:, :, :3])
    cond = np.greater_equal(y[:, :, 0], pixel_threshold)
    activation_pixels = np.where(cond)
    quad_scores, quad_after_nms = nms(y, activation_pixels)
    with Image.open(img_path) as im:
        d_wight, d_height = resize_image(im, cfg.max_predict_img_size)
        scale_ratio_w = d_wight / im.width
        scale_ratio_h = d_height / im.height
        im = im.resize((d_wight, d_height), Image.NEAREST).convert('RGB')
        quad_im = im.copy()
        draw = ImageDraw.Draw(im)
        for i, j in zip(activation_pixels[0], activation_pixels[1]):
            px = (j + 0.5) * cfg.pixel_size
            py = (i + 0.5) * cfg.pixel_size
            line_width, line_color = 1, 'red'
            if y[i, j, 1] >= cfg.side_vertex_pixel_threshold:
                if y[i, j, 2] < cfg.trunc_threshold:
                    line_width, line_color = 2, 'yellow'
                elif y[i, j, 2] >= 1 - cfg.trunc_threshold:
                    line_width, line_color = 2, 'green'
            draw.line([(px - 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size),
                       (px + 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size),
                       (px + 0.5 * cfg.pixel_size, py + 0.5 * cfg.pixel_size),
                       (px - 0.5 * cfg.pixel_size, py + 0.5 * cfg.pixel_size),
                       (px - 0.5 * cfg.pixel_size, py - 0.5 * cfg.pixel_size)],
                      width=line_width, fill=line_color)
        im.save(img_path + '_act.jpg')
        quad_draw = ImageDraw.Draw(quad_im)
        txt_items = []
        for score, geo, s in zip(quad_scores, quad_after_nms, 
                                 range(len(quad_scores))):
            if np.amin(score) > 0:
                quad_draw.line([tuple(geo[0]),
                                tuple(geo[1]),
                                tuple(geo[2]),
                                tuple(geo[3]),
                                tuple(geo[0])], width=2, fill='red')
                rescaled_geo = geo / [scale_ratio_w, scale_ratio_h]
                rescaled_geo_list = np.reshape(rescaled_geo, (8,)).tolist()
                txt_item = ','.join(map(str, rescaled_geo_list))
                txt_items.append(txt_item + '
')
            elif not quiet:
                print('quad invalid with vertex num less then 4.')
        quad_im.save(img_path + '_predict.jpg')
        if len(txt_items) > 0:
            with open(img_path[:-4] + '.txt', 'w') as f_txt:
                f_txt.writelines(txt_items)


def predict_txt(img_path, model_path, pixel_threshold=cfg.pixel_threshold, quiet=False):
    east_detect.load_weights(model_path)
    img = image.load_img(img_path)
    d_wight, d_height = resize_image(img, cfg.max_predict_img_size)
    scale_ratio_w = d_wight / img.width
    scale_ratio_h = d_height / img.height
    img = img.resize((d_wight, d_height), Image.NEAREST).convert('RGB')
    img = image.img_to_array(img)
    img = preprocess_input(img, mode='tf')
    x = np.expand_dims(img, axis=0)
    y = east_detect.predict(x)

    y = np.squeeze(y, axis=0)
    y[:, :, :3] = sigmoid(y[:, :, :3])
    cond = np.greater_equal(y[:, :, 0], pixel_threshold)
    activation_pixels = np.where(cond)
    quad_scores, quad_after_nms = nms(y, activation_pixels)

    txt_items = []
    for score, geo in zip(quad_scores, quad_after_nms):
        if np.amin(score) > 0:
            rescaled_geo = geo / [scale_ratio_w, scale_ratio_h]
            rescaled_geo_list = np.reshape(rescaled_geo, (8,)).tolist()
            txt_items.append(rescaled_geo_list)
        elif not quiet:
            print('quad invalid with vertex num less then 4.')
    return txt_items

# Training
NUM_CLASSES = len('0123456789') + 1  # extra one for 'blank' in CTC loss
MAX_LABEL_LENGTH = 26
DOWNSAMPLE_FACTOR = 4
VALIDATION_RATIO = 0.2

MODEL_OUT_DIR = 'model/'
IMG_SIZE = (256, 32)
BATCH_SIZE = 32  # recommend
EPOCH = 10

# CODE
DECODE_DICT = {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: '_'}  
ENCODE_DICT = {'0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8, '9': 9, '_': 10}