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

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 ......")

你可能感兴趣的:(python程序)