caffe python lmdb讀寫

caffe中可以采取lmdb健值數據庫的方式向網絡中輸入數據。
所以操作lmdb就圍繞“鍵-值“的方式訪問數據庫就好了。

Write

我們可以采用cv2來讀入自己的圖像數據，采用datum格式來存儲數據。
Datum is a Google Protobuf Message class used to store data and optionally a label. A Datum can be thought of a as a matrix with three dimensions: width, height, and channel.

import lmdb
import numpy as np
import cv2
import caffe
from caffe.proto import caffe_pb2

#basic setting
# 這個設置用來存放lmdb數據的目錄
lmdb_file = 'lmdb_data'
batch_size = 256

# create the lmdb file
# map_size指的是數據庫的最大容量，根據需求設置
lmdb_env = lmdb.open(lmdb_file, map_size=int(1e12))
lmdb_txn = lmdb_env.begin(write=True)
# 因為caffe中經常采用datum這種數據結構存儲數據
datum = caffe_pb2.Datum()

item_id = -1
for x in range(1000):
    item_id += 1

    #prepare the data and label
    
    #data = np.ones((3,64,64), np.uint8) * (item_id%128 + 64) #CxHxW array, uint8 or float
    # pic_path設置成圖像目錄, 0表示讀入灰度圖
	data = cv2.imread(pic_path, 0)
    # label 設置圖像的label就行
    label = item_id%128 + 64

    # save in datum
    datum = caffe.io.array_to_datum(data, label)
    keystr = '{:0>8d}'.format(item_id)
    lmdb_txn.put( keystr, datum.SerializeToString() )

    # write batch
    if(item_id + 1) % batch_size == 0:
        lmdb_txn.commit()
        lmdb_txn = lmdb_env.begin(write=True)
        print (item_id + 1)

# write last batch
if (item_id+1) % batch_size != 0:
    lmdb_txn.commit()
    print 'last batch'
    print (item_id + 1)

Read

import caffe
import lmdb
import numpy as np
import cv2
from caffe.proto import caffe_pb2

lmdb_env = lmdb.open('lmdb_data')
lmdb_txn = lmdb_env.begin()
lmdb_cursor = lmdb_txn.cursor()
datum = caffe_pb2.Datum()

for key, value in lmdb_cursor:
    datum.ParseFromString(value)

    label = datum.label
    data = caffe.io.datum_to_array(datum)

    #CxHxW to HxWxC in cv2
    image = np.transpose(data, (1,2,0))
    cv2.imshow('cv2', image)
    cv2.waitKey(1)
    print('{},{}'.format(key, label))