gen_lmdb.py 文件源码

def gen_output(lmdbname, file_list):
    X = np.zeros((len(file_list), 1, HEIGHT, WIDTH), dtype=np.uint8)
    map_size = X.nbytes * 3

    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, "SegmentationClass", i + ".png")
            m = deepcopy(np.asarray(Image.open(filename)))
            for x in range(m.shape[0]):
                for y in range(m.shape[1]):
                    if m[x][y] == 255:
                        m[x][y] = 0
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = 1
            datum.height = m.shape[0]
            datum.width = m.shape[1]
            datum.data = m.tobytes()
            str_id = i
            txn.put(str_id.encode("ascii"), datum.SerializeToString())
            count += 1