如何使用caffe为FCN创建自己的数据集?

时间:2017-10-16 12:01:47

标签: dataset caffe lmdb

如何使用caffe将图像转换为lmdb for fcn?你知道吗,用caffe很容易创建自己的图像分类数据集,但是如何为fcn创建语义段的自己的数据集?

1 个答案:

答案 0 :(得分:1)

使用此代码。进行必要的路径更改。请在使用前仔细阅读代码。

import caffe
import lmdb
from PIL import Image
import numpy as np
import glob
from random import shuffle

# Initialize the Image set:

NumberTrain = 1111 # Number of Training Images

NumberTest = 1112 # Number of Testing Images

Rheight = 380 # Required Height

Rwidth = 500 # Required Width

LabelHeight = 380 # Downscaled height of the label

LabelWidth = 500 # Downscaled width of the label


# Read the files in the Data Folder

inputs_data_train = sorted(glob.glob("/home/<user>/caffe-with_crop/examples/fcn-32s/TrainData/*.jpg"))
inputs_data_valid = sorted(glob.glob("/home/<user>/caffe-with_crop/examples/fcn-32s/ValData/*.jpg"))
inputs_label = sorted(glob.glob("/home/<user>/caffe-with_crop/examples/fcn-32s/VOC2011/SegmentationClass/*.png"))

shuffle(inputs_data_train) # Shuffle the DataSet
shuffle(inputs_data_valid) # Shuffle the DataSet

inputs_Train = inputs_data_train[:NumberTrain] # Extract the training data from the complete set

inputs_Test = inputs_data_valid[:NumberTest] # Extract the testing data from the complete set

# Creating LMDB for Training Data

print("Creating Training Data LMDB File ..... ")

in_db = lmdb.open('/home/<user>/caffe-with_crop/examples/fcn-32s/TrainVOC_Data_lmdb',map_size=int(1e14))

with in_db.begin(write=True) as in_txn:

    for in_idx, in_ in enumerate(inputs_Train):
        print in_idx
        im = np.array(Image.open(in_)) # or load whatever ndarray you need
        Dtype = im.dtype
        im = im[:,:,::-1]
        im = Image.fromarray(im)
        im = im.resize([Rheight, Rwidth], Image.ANTIALIAS)
        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()

# Creating LMDB for Training Labels

print("Creating Training Label LMDB File ..... ")

in_db = lmdb.open('/home/<user>/caffe-with_crop/examples/fcn-32s/TrainVOC_Label_lmdb',map_size=int(1e14))

with in_db.begin(write=True) as in_txn:

    for in_idx, in_ in enumerate(inputs_Train):
        print in_idx
        in_label = in_[:-25]+'VOC2011/SegmentationClass/'+in_[-15:-3]+'png' # Change the numbers as per requirement
        L = np.array(Image.open(in_)) # or load whatever ndarray you need
        Dtype = L.dtype
        L = L[:,:,::-1]
        Limg = Image.fromarray(L)
        Limg = Limg.resize([LabelHeight, LabelWidth],Image.NEAREST) # To resize the Label file to the required size 

        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_db.close()

# Creating LMDB for Testing Data

print("Creating Testing Data LMDB File ..... ")

in_db = lmdb.open('/home/<user>/caffe-with_crop/examples/fcn-32s/TestVOC_Data_lmdb',map_size=int(1e14))

with in_db.begin(write=True) as in_txn:

    for in_idx, in_ in enumerate(inputs_Test):
        print in_idx    
        im = np.array(Image.open(in_)) # or load whatever ndarray you need
        Dtype = im.dtype
        im = im[:,:,::-1]
        im = Image.fromarray(im)
        im = im.resize([Rheight, Rwidth], Image.ANTIALIAS)
        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()

# Creating LMDB for Testing Labels

print("Creating Testing Label LMDB File ..... ")

in_db = lmdb.open('/home/<user>/caffe-with_crop/examples/fcn-32s/TestVOC_Label_lmdb',map_size=int(1e14))

with in_db.begin(write=True) as in_txn:

    for in_idx, in_ in enumerate(inputs_Test):
        print in_idx    
        in_label = in_[:-25]+'VOC2011/SegmentationClass/'+in_[-15:-3]+'png' # Change the numbers as per requirement
        L = np.array(Image.open(in_)) # or load whatever ndarray you need
        Dtype = L.dtype
        L = L[:,:,::-1]
        Limg = Image.fromarray(L)
        Limg = Limg.resize([LabelHeight, LabelWidth],Image.NEAREST) # To resize the Label file to the required size 

        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_db.close()