Создание формата HDF5 для задачи сегментации изображения - PullRequest
Новая
       33

Создание формата HDF5 для задачи сегментации изображения

0 голосов
/ 29 апреля 2018

Я начал писать код на Python для создания HDF5 для задач сегментации изображений. Я использовал код в эту ссылку и ссылку, предоставленную Shai . мои изображения - один канал и в формате .mat. Я написал следующий код, я только хочу проверить с экспертами, является ли этот код правильным или нет. Могут ли эксперты взглянуть? Спасибо 

import os, h5py
import caffe
import numpy as np
import scipy
import scipy.io as sio
from array import array
import cv2
import matplotlib.pyplot as plt

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

def img_to_hdf5(paths_src_file,paths_lbl_file,path_dst,msg):
    """
    paths_src_file : path to the image paths in a txt file
    paths_lbl_file : path to the image paths in a txt file
    path_dst = path to the hdf5 file  
    """

    print(msg)
    arrays = {}
    SIZE=256   #fixed size of all images

    #read the lines of img and lbl path from text file and save into paths_src and paths_lbl
    paths_src = []
    with open(paths_src_file) as f:
        for line in f.readlines():
            line = line.strip('\n')
            paths_src.append(line)

    paths_lbl=[]
    with open(paths_lbl_file) as f:
        for line in f.readlines():
            line=line.strip('\n')
            paths_lbl.append(line)

    data = np.zeros( (len(paths_src), 1, SIZE, SIZE), dtype='f4' )  #1 channel grayscale image
    label = np.zeros( (len(paths_lbl), 1, SIZE, SIZE), dtype='f4' ) #1 channel label image
    for in_idx, in_ in enumerate(paths_src):
        print in_idx,in_
        f=h5py.File(in_,'r')
        mat=f['image'].value
        im=np.array(mat,dtype=np.float32)
        #im = cv2.cvtColor(im,cv2.COLOR_GRAY2RGB)
        #im = im[:,:,::-1]  #switch from RGB to BGR
        im = im.reshape(im.shape[0], im.shape[1], 1)
        im = im.transpose((2,0,1)) # convert to CxHxW
        data[in_idx]=im

    for in_idx, in_ in enumerate(paths_lbl):
        print in_idx,in_
        f=h5py.File(in_,'r')
        mat=f['image'].value
        im=np.array(mat,dtype=np.float32)
        #im = cv2.cvtColor(im,cv2.COLOR_GRAY2RGB)
        #im = im[:,:,::-1]  #switch from RGB to BGR

        im = im.reshape(im.shape[0], im.shape[1], 1)
        im = im.transpose((2,0,1)) # convert to CxHxW
        label[in_idx]=im
    h5_train = os.path.join(path_dst, 'train_data.h5')
    with h5py.File(h5_train,'w') as H:
        H.create_dataset( 'data', data=data ) # note the name X given to the dataset!
        H.create_dataset( 'label', data=label ) # note the name y given to the dataset!
    text_train = os.path.join(path_dst, 'train-path.txt')
    with open(text_train,'w') as L:
        L.write(h5_train) # list all h5 files you are going to use


train_img_paths = './train_img.txt'  #text file of paths to images
train_label_paths = './train_label.txt'  #text file of paths to label images (ground truth)
train_img_hdf5 = '/home/ss/workspace/create_hdf5/' # Your path to h5 file                       

st='Creating Training Data HDF5 File .....'
img_to_hdf5(train_img_paths, train_label_paths,train_img_hdf5,st)
print('DONE...')
...