IAM

OPENSOURCEFAN STUDYING
STUDYINGCOMPUTERSCIENCEANDMATH COMPUTERSCIENCE

Check out the latest superpixel benchmark — Superpixel Benchmark (2016) — and let me know your opinion! @david_stutz
14thJANUARY2017

SNIPPET

Read and write LMDBs using Caffe and Python. The methods read and write are compatible with Caffe’s LMDB format as long as the images are read using OpenCV (i.e. in BGR) format. Integer labels as well as 8-bit and float images are supported.

lmdb_io.py
# -*- coding: utf-8 -*-
"""
Example for reading/writing LMDB.

@author: david
"""

import shutil
import caffe
import numpy
import lmdb
import os

def write (lmdb_path, images, labels):
    """
    Write to the LMDB.
    
    :param lmdb_path: path to LMDB
    :type lmdb_path: string
    :param images: list of images as numpy array
    :type images: [numpy.ndarray]
    :param labels: list of labels
    :type labels: [int]
    """
    
    if len(labels) > 0:
        assert len(images) == len(labels)
        
    env = lmdb.open(lmdb_path, map_size = max(10485760, len(images)*images[0].nbytes))
    
    with env.begin(write = True) as transaction:
        for i in range(len(images)):
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = images[i].shape[2]
            datum.height = images[i].shape[0]
            datum.width = images[i].shape[1]
            
            assert images[i].dtype == numpy.uint8 or images[i].dtype == numpy.float, "currently only numpy.uint8 and numpy.float images are supported"
            
            if images[i].dtype == numpy.uint8:
                datum.data = images[i].transpose(2, 0, 1).tostring()
            else:
                datum.float_data.extend(images[i].transpose(2, 0, 1).flat)
                
            if len(labels) > 0:
                datum.label = labels[i]
            
            key = '{:08}'.format(i)
            transaction.put(key.encode('ascii'), datum.SerializeToString());
        
def read(lmdb_path):
    """
    Read the LMDB.
    
    :param lmdb_path: path to LMDB
    :type lmdb_path: string
    :return: images and labels
    :rtype: ([numpy.ndarray], [int])
    """
    
    images = []
    labels = []
    env = lmdb.open(lmdb_path, readonly = True)
    
    with env.begin() as transaction:
        cursor = transaction.cursor();
        
        for key, raw in cursor:
            datum = caffe.proto.caffe_pb2.Datum()
            datum.ParseFromString(raw)
            
            label = datum.label
            
            if datum.data:
                image = numpy.fromstring(datum.data, dtype = numpy.uint8).reshape(datum.channels, datum.height, datum.width).transpose(1, 2, 0)
            else:
                image = numpy.array(datum.float_data).astype(numpy.float).reshape(datum.channels, datum.height, datum.width).transpose(1, 2, 0)
            
            images.append(image)
            labels.append(label)
    
    return images, labels

if __name__ == '__main__':
    
    images = []
    labels = []
    
    for i in range(10):
        image = numpy.zeros((2, 2, 1)).astype(numpy.uint8)
        label = 0
        
        images.append(image)
        labels.append(label)
    
    lmdb_path = 'test_lmdb'
    if os.path.exists(lmdb_path):
        shutil.rmtree(lmdb_path)
    
    write(lmdb_path, images, labels)
    read_images, read_labels = read(lmdb_path)
    
    print read_images, read_labels

What is your opinion on the code snippet? Is it working? Let me know your thoughts in the comments below or using the following platforms: