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

最新推荐文章于 2024-07-25 22:29:57 发布

menggesijiu

最新推荐文章于 2024-07-25 22:29:57 发布

阅读量1.2k

点赞数 1

分类专栏：深度学习 caffe 文章标签： caffe hdf5格式数据格式转换 Python

本文链接：https://blog.csdn.net/qq_20317319/article/details/70238708

版权

深度学习同时被 2 个专栏收录

2 篇文章 0 订阅

订阅专栏

caffe

1 篇文章 0 订阅

订阅专栏

网上相关资料很少，今天做了将文件夹下的所有.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()