create_lmdb.py 文件源码

def create_image_lmdb(target, samples, bgr=BGR, normalize=False):
    """Create an image LMDB

    Args:
        target (str): path of the LMDB to be created
        samples (array list): list of images to be included in the LMDB
        bgr (bool): True if we want to reverse the channel order (RGB->BGR)
        normalize (bool): True if we want to normalize data in [0,1]
    """

    # Open the LMDB
    if os.path.isdir(target):
        raise Exception("LMDB already exists in {}, aborted.".format(target))
    db = lmdb.open(target, map_size=int(1e12))
    with db.begin(write=True) as txn:
        for idx, sample in tqdm(enumerate(samples), total=len(samples)):
            sample = io.imread(sample)
            # load image:
            if normalize:
                # - in [0,1.]float range
                sample = img_as_float(sample)
            if bgr:
                # - in BGR (reverse from RGB)
                sample = sample[:,:,::-1]
            # - in Channel x Height x Width order (switch from H x W x C)
            sample = sample.transpose((2,0,1))
            datum = caffe.io.array_to_datum(sample)
            # Write the data into the db
            txn.put('{:0>10d}'.format(idx), datum.SerializeToString())

    db.close()