caffe实战(二):汉字识别----------中文文字数据集

2018-07-02  本文已影响0人  黑桃Q_6c7d

https://www.jianshu.com/writer#/notebooks/24210100/notes/28352164 上篇我们介绍了caffe环境的搭建,这次我们开始真正的caffe实战。
模型的训练当然需要一个好的文字集,我们在网上找了好久都没有找到适合的印刷体汉字数据集,于是,我们决定自己生成用于训练和测试的汉字数据集。
一、汉字收集
这要看你需要识别什么汉字了,由于比赛需要,我们收集了常见的3500个汉字以及一些繁体字。如图:

汉字.JPG
二、收集需要用到的字体文件
考虑到我们遇到的文字不一定是单一的楷书或宋体,我们又进一步收集了多种字体文件,来生成不同字体的汉字数据集。以下是我们收集的9种字体文件,包括仿宋、黑体、斜体等等:
字体.JPG
三、生成字体图像,存储在规定的目录下
首先是定义好输入参数,其中包括输出目录、字体目录、测试集大小、图像尺寸等等。以下为部分源码:
description = '''
deep_ocr_make_caffe_dataset --out_caffe_dir /root/data/caffe_dataset
--font_dir /root/workspace/deep_ocr_fonts/chinese_fonts
--width 30 --height 30 --margin 4 --langs lower_eng
'''
parser = argparse.ArgumentParser(
    description=description, formatter_class=RawTextHelpFormatter)
parser.add_argument('--out_caffe_dir', dest='out_caffe_dir',
                    default=None, required=True,
                    help='write a caffe dir')
parser.add_argument('--font_dir', dest='font_dir',
                    default=None, required=True,
                    help='font dir to to produce images')
parser.add_argument('--test_ratio', dest='test_ratio',
                    default=0.3, required=False,
                    help='test dataset size')
parser.add_argument('--width', dest='width',
                    default=None, required=True,
                    help='width')
parser.add_argument('--height', dest='height',
                    default=None, required=True,
                    help='height')
parser.add_argument('--no_crop', dest='no_crop',
                    default=True, required=False,
                    help='', action='store_true')
parser.add_argument('--margin', dest='margin',
                    default=0, required=False,
                    help='', )
parser.add_argument('--langs', dest='langs',
                    default="chi_sim", required=True,
                    help='deep_ocr.langs.*, e.g. chi_sim, chi_tra, digits...')
options = parser.parse_args()

out_caffe_dir = os.path.expanduser(options.out_caffe_dir)
font_dir = os.path.expanduser(options.font_dir)
test_ratio = float(options.test_ratio)
width = int(options.width)
height = int(options.height)
need_crop = not options.no_crop
margin = int(options.margin)
langs = options.langs

image_dir_name = "images"

images_dir = os.path.join(out_caffe_dir, image_dir_name)

图像调整以及汉字图片的生成。我们根据第一步生成的汉字列表来生成对应的汉字图片。


汉字列表.JPG

我们使用的工具是PIL,PIL里面有很好用的汉字生成函数,我们用这个函数再结合我们提供的字体文件,就可以生成我们想要的数字化的汉字了。我们先设定好我们生成的字体颜色为黑底白色,字体尺寸由输入参数来动态设定。


image.png

以下为部分源码;

class Font2Image(object):

def __init__(self,
             width, height,
             need_crop, margin):
    self.width = width
    self.height = height
    self.need_crop = need_crop
    self.margin = margin
def do(self, font_path, char, path_img):
    find_image_bbox = FindImageBBox()
    img = Image.new("RGB", (self.width, self.height), "black")
    draw = ImageDraw.Draw(img)
    font = ImageFont.truetype(font_path, int(self.width * 0.7),)
    draw.text((0, 0), char, (255, 255, 255),
              font=font)
    data = list(img.getdata())
    sum_val = 0
    for i_data in data:
        sum_val += sum(i_data)
    if sum_val > 2:
        np_img = np.asarray(data, dtype='uint8')
        np_img = np_img[:, 0]
        np_img = np_img.reshape((self.height, self.width))
        cropped_box = find_image_bbox.do(np_img)
        left, upper, right, lower = cropped_box
        np_img = np_img[upper: lower + 1, left: right + 1]
        if not self.need_crop:
            preprocess_resize_keep_ratio_fill_bg = \
                PreprocessResizeKeepRatioFillBG(self.width, self.height,
                                                fill_bg=False,
                                                margin=self.margin)
            np_img = preprocess_resize_keep_ratio_fill_bg.do(
                np_img)
        cv2.imwrite(path_img, np_img)
    else:
        print("%s doesn't exist." % path_img)

这里我们把生成的数据集分成两部分,其中任意70%的数据用作训练集,其余30%用作测试集。

 parser.add_argument('--test_ratio', dest='test_ratio',
                    default=0.3, required=False,
                    help='test dataset size')

————————————————————————————————

train_list = []
test_list = []
max_train_i = int(len(verified_font_paths) * (1.0 - test_ratio))
for i, verified_font_path in enumerate(verified_font_paths):
    is_train = True
    if i >= max_train_i:
        is_train = False
    for j, char in enumerate(lang_chars):
        if j not in y_to_tag:
            y_to_tag[j] = char
        char_dir = os.path.join(images_dir, "%d" % j)
        if not os.path.isdir(char_dir):
            os.makedirs(char_dir)
        path_image = os.path.join(
            char_dir,
            "%d_%s.jpg" % (i, os.path.basename(verified_font_path)))
        relative_path_image = os.path.join(
            image_dir_name, "%d"%j, 
            "%d_%s.jpg" % (i, os.path.basename(verified_font_path))
        )
        font2image.do(verified_font_path, char, path_image)
        if is_train:
            train_list.append((relative_path_image, j))
        else:
            test_list.append((relative_path_image, j))

整个代码运行下来,我们生成了五个文件,其中包括一个文字列表还有文字列表对应的json文件(注:这个json文件在后面的caffe模型的使用过程中会用到),一个训练集,一个测试集和每个字对应不同字体的图片集。如下图:


生成文件.JPG

每个字对应不同字体的图片集如下:

字符图片.JPG

至此,数据集生成完毕。
未完,待续>>>

上一篇下一篇

猜你喜欢

热点阅读