python3 制作CRNN lmdb格式的数据集

写入的代码

#-*- coding:utf-8 -*-
 
import os
import lmdb#先pip install这个模块哦
import cv2
import glob
import numpy as np
 
 
def checkImageIsValid(imageBin):
    if imageBin is None:
        return False
    # imageBuf = np.fromstring(imageBin, dtype=np.uint8)
    imageBuf = np.frombuffer(imageBin, dtype=np.uint8)
    # 从网络读取图像数据并转换成图片格式
    # 读入灰度图片，可用0作为实参替代
    img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
    if img is None:
        return False
    imgH, imgW = img.shape[0], img.shape[1]
    if imgH * imgW == 0:
        return False
    return True
 
def writeCache(env, cache):
    with env.begin(write=True) as txn:
        for k, v in cache.items():
            txn.put(str(k).encode(), str(v).encode())
 
def createDataset(outputPath, imagePathList, labelList, lexiconList=None, checkValid=True):
    """
    Create LMDB dataset for CRNN training.
#    ARGS:
        outputPath    : LMDB output path
        imagePathList : list of image path
        labelList     : list of corresponding groundtruth texts
        lexiconList   : (optional) list of lexicon lists
        checkValid    : if true, check the validity of every image
    """
    # print (len(imagePathList) , len(labelList))
    assert(len(imagePathList) == len(labelList))
    nSamples = len(imagePathList)
    print ('...................')
    env = lmdb.open(outputPath, map_size=1099511627776)#1099511627776)所需要的磁盘空间的最小值，之前是1T，我改成了8g，否则会报磁盘空间不足，这个数字是字节
    
    cache = {}
    cnt = 1
    for i in range(nSamples):
        imagePath = imagePathList[i]
        label = labelList[i]
        if not os.path.exists(imagePath):
            print('%s does not exist' % imagePath)
            continue
        with open(imagePath, 'rb') as f:
            imageBin = f.read()

        if checkValid:
            if not checkImageIsValid(imageBin):
                print('%s is not a valid image' % imagePath)#注意一定要在linux下，否则f.read就不可用了，就会输出这个信息
                continue
 
        imageKey = 'image-%09d' % cnt
        labelKey = 'label-%09d' % cnt
        cache[imageKey] = imageBin
        cache[labelKey] = label
        if lexiconList:
            lexiconKey = 'lexicon-%09d' % cnt
            cache[lexiconKey] = ' '.join(lexiconList[i])
        if cnt % 1000 == 0:
            writeCache(env, cache)
            cache = {}
            print('Written %d / %d' % (cnt, nSamples))
        cnt += 1
    nSamples = cnt - 1
    cache['num-samples'] = str(nSamples)
    writeCache(env, cache)
    print('Created dataset with %d samples' % nSamples)
 
 
def read_text(path):
    
    with open(path) as f:
        text = f.read()
    text = text.strip()
    
    return text
 
 
if __name__ == '__main__':

    mode = 'train'
    # mode = 'test'
    # lmdb 输出目录
    outputPath = f'/media/disk3/yrq/ocr/DB/datasets/wm_lmdb/{mode}'#训练集和验证集要跑两遍这个程序，分两次生成
 
    path = f"/media/disk3/yrq/ocr/DB/datasets/water_meter/{mode}_images/*.jpg"
    imagePathList = sorted(glob.glob(path))
    path_label = f"/media/disk3/yrq/ocr/DB/datasets/water_meter/{mode}_gts/*.txt"
    imagelabelList = sorted(glob.glob(path_label))
    print ('------------',len(imagePathList),'------------')
    imgLabelLists = []
    for p,label in zip(imagePathList,imagelabelList):
        try:
           imgLabelLists.append((p, read_text(label)))
        except:
            continue
            
    imgLabelList = sorted(imgLabelLists, key = lambda x:len(x[1]))
    imgPaths = [ p[0] for p in imgLabelList]
    txtLists = [ p[1] for p in imgLabelList]
    
    createDataset(outputPath, imgPaths, txtLists, lexiconList=None, checkValid=True)
 

读取的代码

# -*- coding: utf-8 -*-
import numpy as np
import lmdb
import cv2

with lmdb.open("/media/disk3/yrq/ocr/DB/datasets/wm_lmdb/train") as env:
    txn = env.begin()
    for key, value in txn.cursor():
        print (key,value)
        imageBuf = np.frombuffer(value, dtype=np.uint8)
        img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
        if img is not None:
            cv2.imshow('image', img)
            cv2.waitKey()
        else:
            print ('This is a label: {}'.format(value))