您的位置:首页 > Web前端

caffe 将三通道或四通道图片转换为lmdb格式,将标签(单通道灰度图)转换为lmdb格式

2016-11-28 11:25 471 查看
import numpy as np
import lmdb
from PIL import Image
import sys

# import caffe module
caffe_root = '/home/user/SegNet/caffe-segnet/'
sys.path.insert(0, caffe_root + 'python')
import caffe

# # read file
train_file = open('/home/user/train.txt')
inputs_data_train = train_file.readlines()
train_file.close()

print("Creating Training Data LMDB File ..... ")
in_db = lmdb.open('/home/user/Val_Data_lmdb', map_size=int(1e12))
with in_db.begin(write=True) as in_txn:
for in_idx, in_ in enumerate(inputs_data_train):
# print in_idx
in_ = in_.strip()
im = np.array(Image.open(in_))
Dtype = im.dtype
if im.shape[2]== 3:
print('The image has 3 channel')
# RGB to BGR
im = im[:, :, ::-1]
if im.shape[2]==4:
im3 = im[:,:,0:3]
im3 = np.array(im3)
#RGB to BGR
im3 = im3[:,:,::-1]
im[:, :, 0] = im3[:, :, 0]
im[:, :, 1] = im3[:, :, 1]
im[:, :, 2] = im3[:, :, 2]
im = Image.fromarray(im)
im = np.array(im, Dtype)
im = im.transpose((2, 0, 1))
im_dat = caffe.io.array_to_datum(im)
in_txn.put('{:0>10d}'.format(in_idx), im_dat.SerializeToString())
in_db.close()

# read file
label_file = open('/home/user/trainannot.txt')
inputs_data_label = label_file.readlines()
label_file.close()

print("Creating Training Label LMDB File ..... ")
#map_size:Change the maximum size of the map file
in_db1 = lmdb.open('//home/user/Val_Label_Data_lmdb',map_size=int(1e12))
with in_db1.begin(write=True) as in_txn:
for in_idx, in_ in enumerate(inputs_data_label):
in_ = in_.strip()
Dtype = 'uint8'
L = np.array(Image.open(in_), Dtype)
Limg = Image.fromarray(L)
L = np.array(Limg,Dtype)
L = L.reshape(L.shape[0],L.shape[1],1)
L = L.transpose((2,0,1))
L_dat = caffe.io.array_to_datum(L)
in_txn.put('{:0>10d}'.format(in_idx),L_dat.SerializeToString())
in_db1.close()

print("Finish creating lmdb file ......")
内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: