利用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)