gen_lmdb.py 文件源码

def make_lmdb_input(lmdbname, channel_directories, range_set):

    X = np.zeros((len(range_set), len(channel_directories), WIDTH, HEIGHT), dtype=np.double)
    map_size = X.nbytes * 10

    env = lmdb.open(lmdbname, map_size=map_size)

    count = 0
    for i in range_set:
        with env.begin(write=True) as txn:
            filename = str(i) + ".png"
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = X.shape[1]
            datum.height = X.shape[2]
            datum.width = X.shape[3]
            for j in range(len(channel_directories)):
                dirname = channel_directories[j]
                X[count][j] = np.asarray(PIL.Image.open(os.path.join(dirname, filename)), dtype=np.double)
            datum.data = X[count].tobytes()
            str_id = '{:08}'.format(count)
            txn.put(str_id.encode("ascii"), datum.SerializeToString())
            count += 1