CTPN 数据处理/dataloader

深度学习模型训练的第一步就是准备数据，制作标签（gt, ground truth）。然后根据gt和预测的值之差通过梯度下降的方法优化模型参数。
CTPN中gt包括两部分，一是分类的gt，二是bbox（检测框）的gt。
下面这个函数是CTPN的数据处理主要函数。tensorflow1中读取数据可以使用多线程读取，因为读取数据是用cpu读取的，为了高效利用GPU，使用多线程读取效率比较高。这一部分比较简单，返回的是图片，bbox以及图片信息（高、宽、通道数）。bbox这里使用的是绝对坐标表示，[x_min, y_min, x_max, y_max, 1]，最后一位1表示这个bbox是文字。有个地方需要注意，这里返回使用的是yield，它的作用和return一样，不同之处在于，yield返回结果之后并函数还会接着运行。

def generator(vis=False):
    image_list = np.array(get_training_data())
    print('{} training imas in {}'.format(image_list.shape[0], DATA_FOLDER))
    index = np.arange(0, image_list.shape[0])
    while True:
        np.random.shuffle(index)
        for i in index:
            print(i)
            try:
                im_fn = image_list[i]
                im = cv2.imread(im_fn)
                h, w, c = im.shape
                im_info = np.array([h, w, c]).reshape([1, 3])

                _, fn = os.path.split(im_fn)
                fn, _ = os.path.splitext(fn)
                txt_fn = os.path.join(DATA_FOLDER, "label", fn + '.txt')
                if not os.path.exists(txt_fn):
                    print("Ground truth for image {} not exist!".format(im_fn))
                    continue
                bbox = load_annoataion(txt_fn)
                if len(bbox) == 0:
                    print("Ground truth for image {} empty!".format(im_fn))
                    continue

                if vis:
                    for p in bbox:
                        cv2.rectangle(im, (p[0], p[1]), (p[2], p[3]), color=(0, 0, 255), thickness=1)
                    fig, axs = plt.subplots(1, 1, figsize=(30, 30))
                    axs.imshow(im[:, :, ::-1])
                    axs.set_xticks([])
                    axs.set_yticks([])
                    plt.tight_layout()
                    plt.show()
                    plt.close()
                yield [im], bbox, im_info

            except Exception as e:
                print(e)
                continue

def load_annoataion(p):
    bbox = []
    with open(p, "r") as f:
        lines = f.readlines()
    for line in lines:
        line = line.strip().split(",")
        x_min, y_min, x_max, y_max = map(int, line)
        bbox.append([x_min, y_min, x_max, y_max, 1])
    return bbox

def get_training_data():
    img_files = []
    exts = ['jpg', 'png', 'jpeg', 'JPG']
    for parent, dirnames, filenames in os.walk(os.path.join(DATA_FOLDER, "image")):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    img_files.append(os.path.join(parent, filename))
                    break
    print('Find {} images'.format(len(img_files)))
    return img_files