利用python生成caffe训练数集的txt文件

利用caffe训练自己的数据，需要生成一个如下所示的txt文件
H1855.jpg 1
W33481.jpg 2
W524.jpg 2
H19114.jpg 1
W5754.jpg 2
H46204.jpg 1
H33060.jpg 1
H41133.jpg 1
其代码如下

#! /usr/bin/env python
#-*- coding: UTF-8 -*- 
import os
import re

#生成图片列表清单txt文件
def createFileList(images_path,txt_path):
    fw = open(txt_path,"wr")
    images_name = os.listdir(images_path)

    #遍历所有文件名
    for eachname in images_name:
        #正则表达式
        Hz = r'(^H\d{0,10}.jpg$)'
        Wz = r'(^W\d{0,10}.jpg$)'

        #正则表达式匹配
        Hz_name = re.search(Hz,eachname)
        Wz_name = re.search(Wz,eachname)

        #按照规则将内容写入txt文件中

        if Hz_name != None:
            fw.write(Hz_name.group(0) + ' 1\n')
        if Wz_name != None:
            fw.write(Wz_name.group(0) + ' 2\n')

    print("生成txt文件成功")

    #关闭fw
    fw.close()

caffe_root = '/home/yh/caffe/'
my_caffe_project = caffe_root + 'examples/H-W/'

#图片存放目录
images_path = caffe_root + 'examples/H-W/train/'
#生成的图片列表清单txt文件名
txt_name = 'trainlist.txt'
#生成的列表清单的保存目录
txt_save_path = my_caffe_project + txt_name

#生成txt文件
createFileList(images_path,txt_save_path)