您的位置:首页 > Web前端

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

2017-07-01 10:22 645 查看
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 ......")
内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: