gen_lmdb.py 文件源码

def gen_input(lmdbname, file_list):
    X = np.zeros((len(file_list), 3, HEIGHT, WIDTH), dtype=np.float32)
    map_size = X.nbytes * 5

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

    count = 0
    for i in file_list:
        print count
        with env.begin(write=True) as txn:
            filename = os.path.join(DIR, "JPEGImages", i + ".jpg")
            m = np.asarray(Image.open(filename)).transpose((2, 0, 1))
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = m.shape[0]
            datum.height = m.shape[1]
            datum.width = m.shape[2]
            datum.data = m.tobytes()
            str_id = i
            txn.put(str_id.encode("ascii"), datum.SerializeToString())
            count += 1