CTPN 數(shù)據(jù)處理/dataloader

深度學(xué)習(xí)模型訓(xùn)練的第一步就是準(zhǔn)備數(shù)據(jù),制作標(biāo)簽(gt, ground truth)。然后根據(jù)gt和預(yù)測(cè)的值之差通過梯度下降的方法優(yōu)化模型參數(shù)。
CTPN中g(shù)t包括兩部分,一是分類的gt,二是bbox(檢測(cè)框)的gt。
下面這個(gè)函數(shù)是CTPN的數(shù)據(jù)處理主要函數(shù)。tensorflow1中讀取數(shù)據(jù)可以使用多線程讀取,因?yàn)樽x取數(shù)據(jù)是用cpu讀取的,為了高效利用GPU,使用多線程讀取效率比較高。這一部分比較簡(jiǎn)單,返回的是圖片,bbox以及圖片信息(高、寬、通道數(shù))。bbox這里使用的是絕對(duì)坐標(biāo)表示,[x_min, y_min, x_max, y_max, 1],最后一位1表示這個(gè)bbox是文字。有個(gè)地方需要注意,這里返回使用的是yield,它的作用和return一樣,不同之處在于,yield返回結(jié)果之后并函數(shù)還會(huì)接著運(yùn)行。

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
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。