caffe中将某个目录下的jpg图像转换为hdf5类型

网上相关资料很少，今天做了将文件夹下的所有.jpg格式的图片转换成hdf5格式，中途踩了很多坑，最后在老师的指导下写出来了，特此记录。

二话不说，直接上代码，很好理解。我的图片命名是x1_y1_x2_y2.jpg的形式，用CNN做回归，找到2个坐标的x,y值。建议用python做文本处理，非常快，很容易上手。

主要看data和label具体怎么写，怎么定义，怎么转换，这个是关键。所有的图片都读取完，一次性create_dataset，不然会出错。

import numpy as np
import h5py
import cv2
import os
import random
import code
import string

def store_hdf5(filename, mapping):
"""Function to store data mapping to a hdf5 file

Args:
filename (str): The output filename
mapping (dic): A dictionary containing mapping from name to numpy data The complete mapping will be stored as single datasets in the h5py file.
print("Storing hdf5 file %s" % filename)
with h5py.File(filename, 'w') as hf:
for label, data in mapping.items():
print("  adding dataset %s with shape %s" % (label, data.shape))
hf.create_dataset(label, data=data)

print("  finished")

if __name__=="__main__":

src_dir = "/usr/MyProject/chineseWordsRecognition/meta/desdir2"
des_dir = "/usr/caffe_20150410/examples/licensefile/train_desdir"
des_txt = "/usr/caffe_20150410/examples/licensefile/train_destxt.txt"

flist = os.listdir(src_dir)

if not os.path.exists(des_dir):
os.mkdir(des_dir)

print "num of files: " + str(len(flist))

random.shuffle(flist)

des_txt_file = open(des_txt,'w')

#max_num=1000000 #56GB
max_num=32000 #one split contains 32000 images
W = 210
H = 30

TOT_MAX=1000000

data = np.zeros((max_num,3, H, W),dtype=np.float32)
labels = np.full((max_num, 4), -1, dtype=np.float32)

CNT=0;

TOT = len(flist)

TOT_CNT = 0
BATCH_CNT=0

for fname in flist:
if fname[0]=="." or fname[0]=="_":
continue
src_path = src_dir + "/" + fname

res = cv2.imread(src_path, cv2.IMREAD_COLOR).astype(np.float32)
tmpimg = cv2.resize(res,(W,H))

h,w,ch = tmpimg.shape
if h==0 or w==0:
continue

img = np.zeros((3,H,W),dtype=np.float32)
img[0,:,:] = tmpimg[:,:,0]-104 #图片归一化，除以每个通道的均值
img[1,:,:] = tmpimg[:,:,1]-117
img[2,:,:] = tmpimg[:,:,2]-123

annot = fname.split(".")[0].split("_")[0:4] #取图片命名的四个坐标

if len(annot)!=4:
continue

resh,resw,resch = res.shape
l = string.atof(annot[0])/resw
labels[CNT,0] = l
l = string.atof(annot[1])/resh   #坐标一定要归一化，切记切记
labels[CNT,1] = l
l = string.atof(annot[2])/resw
labels[CNT,2] = l
l = string.atof(annot[3])/resh
labels[CNT,3] = l

data[CNT,:,:,:] = img

CNT+=1
TOT_CNT+=1
if TOT_CNT == 10000:   #hdf5文件的大小有限制，如果太大的话，一定要分批写
break

if TOT_CNT%1000==0:
print "TOT=" + str(TOT) + ", TOT_CNT=" + str(TOT_CNT)
print code

if CNT>=max_num:
des_path = des_dir + "/" + str(BATCH_CNT) + ".h5"
des_txt_file.write(des_path + "\n")
CNT=0
BATCH_CNT+=1
store_hdf5(des_path, {"data" : data, "labels" : labels})
data = np.zeros((max_num,3, H, W),dtype=np.float32)
labels = np.full((max_num, 4), -1, dtype=np.float32)

if TOT_CNT>=TOT_MAX-1:
break

if CNT>0:
data = data[0:CNT, :, :, :]

labels = labels[0:CNT, :]

des_path = des_dir + "/" + str(BATCH_CNT) + ".h5"
des_txt_file.write(des_path + "\n")
store_hdf5(des_path, {"data" : data, "labels" : labels})
des_txt_file.close()