LMDB

1. Write image to LMDB:

#!/usr/bin/python3
import numpy as np
import lmdb
import caffe
import cv2
import matplotlib.pyplot as plt

# set max image number
N = 1000

# prepare a image which will be saved into lmdb
img = cv2.imread("travel.jpg")
img = cv2.resize(img, (64,48))
plt.figure("org")
plt.imshow(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
plt.show()

# swap the axes to fit caffe
img = np.swapaxes(img, 1, 2)
img = np.swapaxes(img, 0, 1)
print("swaped image shape: ", img.shape)

# the order of the axes in caffe used to be (N d h w)
# create a buffer named X
X = np.zeros((N, 3, 48, 64), dtype=np.uint8) 
X[0] = img
y = np.zeros(N, dtype=np.int64)

# data size
map_size = X.nbytes*10

# dump data into lmdb
with lmdb.open("mylmdb", map_size=map_size) as env:
    with env.begin(write=True) as txn:
        for i in range(N):
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = X.shape[1] 
            datum.height = X.shape[2]
            datum.width = X.shape[3]

            datum.data = X[i].tobytes()
            datum.label = int(y[i])

            id = "{:08}".format(i)
            txn.put(id.encode(), datum.SerializeToString())

2. Read image from LMDB:

#!/usr/bin/python3
import numpy as np
import lmdb
import caffe
import cv2
import matplotlib.pyplot as plt

with lmdb.open("mylmdb") as env:
    with env.begin() as txn:
        with txn.cursor() as cursor:
            datum = caffe.proto.caffe_pb2.Datum()
            i=0
            for key, value in cursor:
                i=i+1
                datum.ParseFromString(value)
                label = datum.label
                data = caffe.io.datum_to_array(datum)
                
                data = np.swapaxes(data, 0, 1)
                data = np.swapaxes(data, 1, 2)
                plt.figure("datum")
                plt.imshow(cv2.cvtColor(data, cv2.COLOR_RGB2BGR))
                plt.show()

                print("Counter:", i)
                print("Data shape: ", data.shape)
                print("Label:", label)
                break

留言

熱門文章