https://www.jianshu.com/writer#/notebooks/24210100/notes/28352164 上篇我们介绍了caffe环境的搭建,这次我们开始真正的caffe实战。
模型的训练当然需要一个好的文字集,我们在网上找了好久都没有找到适合的印刷体汉字数据集,于是,我们决定自己生成用于训练和测试的汉字数据集。
一、汉字收集
这要看你需要识别什么汉字了,由于比赛需要,我们收集了常见的3500个汉字以及一些繁体字。如图:
二、收集需要用到的字体文件
考虑到我们遇到的文字不一定是单一的楷书或宋体,我们又进一步收集了多种字体文件,来生成不同字体的汉字数据集。以下是我们收集的9种字体文件,包括仿宋、黑体、斜体等等:
三、生成字体图像,存储在规定的目录下
首先是定义好输入参数,其中包括输出目录、字体目录、测试集大小、图像尺寸等等。以下为部分源码:
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)
图像调整以及汉字图片的生成。我们根据第一步生成的汉字列表来生成对应的汉字图片。
我们使用的工具是PIL,PIL里面有很好用的汉字生成函数,我们用这个函数再结合我们提供的字体文件,就可以生成我们想要的数字化的汉字了。我们先设定好我们生成的字体颜色为黑底白色,字体尺寸由输入参数来动态设定。
以下为部分源码;
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模型的使用过程中会用到),一个训练集,一个测试集和每个字对应不同字体的图片集。如下图:
每个字对应不同字体的图片集如下:
至此,数据集生成完毕。
未完,待续>>>