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