caffe专题二将图像写入LMDB与读取LMDB—python实现

转自:https://blog.csdn.net/langb2014/article/details/52995349    大牛的文章啊,学习了

调用的是caffe的python接口

一:将图像写入LMDB

import os
import glob
import random
import numpy as np

import cv2

import caffe
from caffe.proto import caffe_pb2
import lmdb

#Size of images
IMAGE_WIDTH = 227
IMAGE_HEIGHT = 227

# train_lmdb、validation_lmdb 路径
train_lmdb = '/xxxxx/train_lmdb'
validation_lmdb = '/xxxxx/validation_lmdb'

# 如果存在了这个文件夹, 先删除
os.system('rm -rf  ' + train_lmdb)
os.system('rm -rf  ' + validation_lmdb)

# 读取图像
train_data = [img for img in glob.glob("/xxx/*jpg")]
test_data = [img for img in glob.glob("/xxxx/*jpg")]

# Shuffle train_data
# 打乱数据的顺序
random.shuffle(train_data)

# 图像的变换, 直方图均衡化, 以及裁剪到 IMAGE_WIDTH x IMAGE_HEIGHT 的大小
def transform_img(img, img_width=IMAGE_WIDTH, img_height=IMAGE_HEIGHT):
    #Histogram Equalization
    img[:, :, 0] = cv2.equalizeHist(img[:, :, 0])
    img[:, :, 1] = cv2.equalizeHist(img[:, :, 1])
    img[:, :, 2] = cv2.equalizeHist(img[:, :, 2])

    #Image Resizing, 三次插值
    img = cv2.resize(img, (img_width, img_height), interpolation = cv2.INTER_CUBIC)
    return img

def make_datum(img, label):
    #image is numpy.ndarray format. BGR instead of RGB
    return caffe_pb2.Datum(
        channels=3,
        width=IMAGE_WIDTH,
        height=IMAGE_HEIGHT,
        label=label,
        data=np.rollaxis(img, 2).tobytes()) # or .tostring() if numpy < 1.9

# 打开 lmdb 环境, 生成一个数据文件,定义最大空间, 1e12 = 1000000000000.0
in_db = lmdb.open(train_lmdb, map_size=int(1e12)) 
with in_db.begin(write=True) as in_txn: # 创建操作数据库句柄
    for in_idx, img_path in enumerate(train_data):
        if in_idx %  6 == 0: # 只处理 5/6 的数据作为训练集
            continue         # 留下 1/6 的数据用作验证集
        # 读取图像. 做直方图均衡化、裁剪操作
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img = transform_img(img, img_width=IMAGE_WIDTH, img_height=IMAGE_HEIGHT)

        if 'cat' in img_path: # 组织 label, 这里是如果文件名称中有 'cat', 标签就是 0
            label = 0         # 如果图像名称中没有 'cat', 有的是 'dog', 标签则为 1
        else:                 # 这里方, label 需要自己去组织
            label = 1         # 每次情况可能不一样, 灵活点

        datum = make_datum(img, label)
        # '{:0>5d}'.format(in_idx):
        #      lmdb的每一个数据都是由键值对构成的,
        #      因此生成一个用递增顺序排列的定长唯一的key
        in_txn.put('{:0>5d}'.format(in_idx), datum.SerializeToString()) #调用句柄,写入内存
        print '{:0>5d}'.format(in_idx) + ':' + img_path

# 结束后记住释放资源,否则下次用的时候打不开。。。
in_db.close() 

# 创建验证集 lmdb 格式文件
print '\nCreating validation_lmdb'
in_db = lmdb.open(validation_lmdb, map_size=int(1e12))
with in_db.begin(write=True) as in_txn:
    for in_idx, img_path in enumerate(train_data):
        if in_idx % 6 != 0:
            continue
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img = transform_img(img, img_width=IMAGE_WIDTH, img_height=IMAGE_HEIGHT)
        if 'cat' in img_path:
            label = 0
        else:
            label = 1
        datum = make_datum(img, label)
        in_txn.put('{:0>5d}'.format(in_idx), datum.SerializeToString())
        print '{:0>5d}'.format(in_idx) + ':' + img_path
in_db.close()
print '\nFinished processing all images'

二:配置文件中使用proto

layer {
  top: "data"
  top: "label"
  name: "data"
  type: "Data"
  data_param {
    source: "/xxxxx/train_lmdb"
    backend:LMDB
    batch_size: 128
  }
  transform_param {
     #mean_file: "/xxxxx/mean.binaryproto"
     mirror: true
  }
  include: { phase: TRAIN }
}

三:读取proto数据

import caffe
from caffe.proto import caffe_pb2

import lmdb
import cv2
import numpy as np

lmdb_env = lmdb.open('mylmdb', readonly=True) # 打开数据文件
lmdb_txn = lmdb_env.begin() # 生成处理句柄
lmdb_cursor = lmdb_txn.cursor() # 生成迭代器指针
datum = caffe_pb2.Datum() # caffe 定义的数据类型

for key, value in lmdb_cursor: # 循环获取数据
    datum.ParseFromString(value) # 从 value 中读取 datum 数据

    label = datum.label
    data = caffe.io.datum_to_array(datum)
    print data.shape
    print datum.channels
    image = data.transpose(1, 2, 0)
    cv2.imshow('cv2.png', image)
    cv2.waitKey(0)

cv2.destroyAllWindows()
lmdb_env.close()

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

Caffe多标签输入常用的的方法有以下几种:
1. 修改Caffe源码使其支持多标签输入,参考CSDN博客
《caffe 实现多标签输入(multilabel、multitask)》
2. HDF5 + Slice Layer,HDF5支持多标签,但Caffe在读取HDF5时会将所有数据一次性预读进内存中,在数据量较大时就对内存有较高要求了;也可以对数据分片,通过在prototxt文件中加入Slice Layer层使Caffe依次读取每个数据分片。
3. 使用两个data输入(例如两个LMDB,一个存储图片,一个存储多标签),然后修改prototxt文件配置两个data layer。

实现第三种方法支持多标签输入。Caffe中单标签LMDB的创建可以通过自带的脚本很方便的实现,但多标签LMDB的创建并不像单标签那样简单易用,需要使用者对Caffe的Data Layer有足够的了解。

多标签的场景下,数据是 N x H x W x C 的一个4维的blob,对应的标签是 N x M x 1 x 1的一个4维的blob。

创建图片LMDB

使用下面的代码创建两个图片LMDB:
- train_data_lmdb - 用于训练的图片LMDB 
- val_data_lmdb - 用于测试的图片LMDB

# 生成训练图片列表文件,即将tarin_images_dir目录下所有图片的路径写入temp.txt文件
find tarin_images_dir -type f -exec echo {} \; > temp.txt

# 在temp.txt文件中每一行后追加伪标签,伪标签起占位符作用,实际并不使用 
sed "s/$/ 0/" temp.txt > train_images.txt

# 根据train_images.txt创建train_data_lmdb
$CAFFE_HOME/build/tools/convert_imageset -resize_height=256 -resize_width=256 / train_images.txt train_data_lmdb

# 计算图片均值
$CAFFE_HOME/build/tools/compute_image_mean train_data_lmdb mean.binaryproto

上面的代码用来生成train_data_lmdb,接下来修改代码生成val_data_lmdb.
创建标签LMDB

使用下面的代码创建两个标签LMDB:
- train_label_lmdb - 用于训练的标签LMDB 
- val_label_lmdb - 用于测试的标签LMDB

读入每幅图片的多标签,生成标签LMDB,注意图片LMDB和标签LMDB中的顺序一致,创建LMDB的部分Python
代码如下:
import sys 
import numpy as np
import lmdb

caffe_root = '/caffe/' 
sys.path.insert(0, caffe_root + 'python')
import caffe

# 根据多标签的位置选择从数据库、文件等中读取每幅图片的多标签,将其构造成一维的np.array类型,并追加入all_labels列表
all_labels = []
# Add your code of reading labels here !

# 创建标签LMDB
key = 0
lmdb_path = "/train_label_lmdb"
env = lmdb.open(lmdb_path, map_size=map_size)
with env.begin(write=True) as txn: 
    for labels in all_labels:
        datum = caffe.proto.caffe_pb2.Datum()
        datum.channels = labels.shape[0]
        datum.height = 1
        datum.width =  1
        datum.data = labels.tostring()          # or .tobytes() if numpy < 1.9 
        datum.label = 0
        key_str = '{:08}'.format(key)

        txn.put(key_str.encode('ascii'), datum.SerializeToString())
        key += 1

上面的代码用来生成train_label_lmdb,接下来修改代码生成val_label_lmdb.

修改prototxt文件

上面我们已经创建了四个数据集train_data_lmdb、val_data_lmdb、train_label_lmdb、val_label_lmdb,并且计算了图片均值mean.binaryproto。

接下来就是修改神经网络deploy.prototxt使其支持多标签训练。由于有四个数据集,我们只需要在deploy.txt中加入四个data layers,分别用来配置这四个数据集。需要注意的是在标签lmdb对应的data layer的参数transform_param中通过scaling参数对label进行缩放,该操作的作用是按照Sigmoid Cross Entropy Loss函数的要求将label范围从[0,255]正则化到[0,1](根据采用的多标签Loss函数的要求,scaling参数可以修改)。

将deploy.prototxt中原有的data layers修改为如下的四个data layers:

# ------------- 配置训练阶段的图片数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "data"
  include {
    phase: TRAIN
  }
  transform_param {
    mirror: true
    crop_size: 224
    mean_file: "data.binaryproto"    # 修改为你的均值文件路径
  }
  data_param {
    source: "train_data_lmdb"        # 修改为你的图片lmdb路径
    batch_size: 32
    backend: LMDB
  }
}

# ------------- 配置训练阶段的标签数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "label"
  include {
    phase: TRAIN
  }
  transform_param {
    scale: 0.00390625                 # 根据需求设置标签缩放系数
    mean_value: 0
  }
  data_param {
    source: "train_label_lmdb"        # 修改为训练集图片lmdb路径
    batch_size: 32
    backend: LMDB
  }
}

# ------------- 配置测试阶段的图片数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "data"
  include {
    phase: TEST
  }
  transform_param {
    mirror: false
    crop_size: 224
    mean_file: "mean.binaryproto"      # 修改为你的均值文件路径
  }
  data_param {
    source: "val_data_lmdb"            # 修改为训练阶段的图片lmdb路径
    batch_size: 1
    backend: LMDB
  }
}

# ------------- 配置测试阶段的标签数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "label"
  include {
    phase: TEST
  }
  transform_param {
    scale: 0.00390625                  # 根据需求设置标签缩放系数
    mean_value: 0
  }
  data_param {
    source: "val_label_lmdb"           # 修改为测试阶段的标签lmdb路径
    batch_size: 1
    backend: LMDB
  }
}

......

# ----------------- 多标签损失函数  -------------------
layers {
  name: "loss"
  type: SIGMOID_CROSS_ENTROPY_LOSS
  bottom: "fc8"                       # 根据需求配置
  bottom: "label"
  top: "loss"
}








猜你喜欢

转载自blog.csdn.net/runner668/article/details/80399275