前言:
系统环境为:win10,本次用的是tensorflow1.12.0
完整目录为:
python人工智能 利用Tensorflow object detection API 搭建物体识别模型(一)CPU版
python人工智能 利用Tensorflow object detection API 搭建物体识别模型(二)CPU版
python人工智能 利用Tensorflow object detection API 搭建物体识别模型(三)CPU版
上篇我们已经成功实现了运行官方demo,这次我们来训练自己的模型
一、下载要训练的素材图片
(1)先创建好目录:D:\ai\objDetect3\mubiaojiance
然后把https://pan.baidu.com/s/1YF1SVRbI4I8JDSXao3RYCQ?pwd=9dew这个素材下载并解压n01440764目录下
(2)用代码筛选出图片宽高符合规定的图片
cd /d D:\ai\objDetect3\mubiaojiance jupyter notebook
然后在自动打开的浏览器中新建一个python文件,命名为get_some_qualified_images
接着输入如下代码
import os import random from PIL import Image import shutil #获取1000张图片中随机选出数量为sample_number*2的一部分图片的路径 def get_some_imagePath(dirPath, sample_number): fileName_list = os.listdir(dirPath) all_filePath_list = [ os.path.join(dirPath, fileName) for fileName in fileName_list ] all_imagePath_list = [ filePath for filePath in all_filePath_list if '.jpg' in filePath ] some_filePath_list = random.sample( all_filePath_list, k=sample_number * 2) return some_filePath_list #获取一部分像素足够,即长,宽都大于300的图片 def get_some_qualified_images(dirPath, sample_number, new_dirPath): some_imagePath_list = get_some_imagePath(dirPath, sample_number) if not os.path.isdir(new_dirPath): os.mkdir(new_dirPath) i = 0 for imagePath in some_imagePath_list: image = Image.open(imagePath) width, height = image.size if width > 300 and height > 300: i += 1 new_imagePath = 'selected_images/%03d.jpg' % i #在处理图像的时候常常需要将一个图像复制到另一个文件夹中,Python可以很方便的利用shutil.copy(src,dst)函数实现这个操作 #返回复制图像的文件路径 shutil.copy( imagePath, new_imagePath) if i == sample_number: break #获取数量为100的合格样本存放到selected_images文件夹中 get_some_qualified_images('n01440764', 100, 'selected_images')
选中这个单元格,点击运行
这时能看到生成的selected_images文件夹,里面总共100张图片,这步操作完成。
(3)为了节省训练时间,我们把筛选出来的这些图片缩小一点。
我们这个时候再插入一段代码,也是选中单元格后运行。
import os from PIL import Image def get_small_images(dirPath, new_dirPath): fileName_list = os.listdir(dirPath) filePath_list = [os.path.join(dirPath, fileName) for fileName in fileName_list] imagePath_list = [filePath for filePath in filePath_list if '.jpg' in filePath] if not os.path.isdir(new_dirPath): os.mkdir(new_dirPath) for imagePath in imagePath_list: image = Image.open( imagePath ) width, height = image.size imageName = imagePath.split('\\')[-1] save_path = os.path.join(new_dirPath, imageName) if width >= 600 and height >= 600: minification = min(width, height) // 300 #缩小倍数 new_width = width // minification new_height = height // minification resized_image = image.resize( (new_width, new_height),Image.ANTIALIAS ) print('图片%s由原来的宽%d,高%d,缩小为宽%d,高%d' % (imageName, width, height, new_width, new_height)) resized_image.save(save_path) else: image.save(save_path) get_small_images('selected_images', 'smaller_images')
运行结果如下:
可以进行下一步了。
二、手动给图片打标签,圈出哪个是鱼
使用打标签工具LabelImg,链接:labelImg-windows_v1.8.1.zip
(1)下载后解压到D:\ai\objDetect3\labelImg-windows_v1.8.1,然后运行labelImg.exe
(2)打开之后,首先点击open dir,将smaller_images文件夹导入进来
(3)此时开始打标签,快捷键是W, 每次打完都需要按下ctrl+s保存一下
键盘快捷键:
A: 上一张图片 D: 下一张图片 W: 开始框选
框选完成后输入fish,表示这是鱼。接着ctrl+s保存即可。直到把100张图片都标记完。
标记完了,然后直接关掉对话框,可以看到smaller_images目录下每张图片都有对应一个xml文件。这个就是我们标记的数据
三、xml数据转换
(1)xml转csv
我们直接从上面jupyter notebook打开的代码框中输入如下代码并运行:
import os import pandas as pd import xml.etree.ElementTree as ET from sklearn.model_selection import train_test_split def xmlPath_list_to_df(xmlPath_list): xmlContent_list = [] for xmlPath in xmlPath_list: print(xmlPath) tree = ET.parse(xmlPath) root = tree.getroot() for member in root.findall('object'): value = ( root.find('filename').text,#文件名 int( root.find('size')[0].text),#width int( root.find('size')[1].text),#height member[0].text,#标签 int( member[4][0].text),#xmin int( member[4][1].text),#ymin int( member[4][2].text),#xmax int( member[4][3].text)#ymax ) xmlContent_list.append(value) column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] xmlContent_df = pd.DataFrame( xmlContent_list, columns = column_name ) return xmlContent_df def dirPath_to_csv(dirPath): fileName_list = os.listdir(dirPath) all_xmlPath_list = [os.path.join(dirPath, fileName) for fileName in fileName_list if '.xml' in fileName] train_xmlPath_list, test_xmlPath_list = train_test_split(all_xmlPath_list, test_size=0.1, random_state=1) train_df = xmlPath_list_to_df( train_xmlPath_list) train_df.to_csv('train.csv') print('成功产生文件train.csv,训练集共有%d张图片' % len(train_xmlPath_list) ) test_df = xmlPath_list_to_df(test_xmlPath_list) test_df.to_csv('test.csv') print('成功产生文件test.csv,测试集共有%d张图片' % len(test_xmlPath_list) ) dirPath_to_csv('smaller_images')
选中单元格,并点击Run运行后输入下面这些东西,说明成功了
将函数train_test_split的参数random_state的值设为1,这样每次划分的训练集和测试集总是相同。如果不设置此参数,则每次划分的训练集和测试集不同。
从上面jupyter notebook打开的代码框中输入如下代码并运行:
#csv转tfrecords import os import pandas as pd import tensorflow as tf from object_detection.utils import dataset_util import shutil def csv2tfrecord( csv_path, imageDir_path, tfrecord_path): objectInfo_df = pd.read_csv(csv_path) tfrecord_writer = tf.python_io.TFRecordWriter(tfrecord_path) for filename, group in objectInfo_df.groupby('filename'): height = group.iloc[0]['height'] width = group.iloc[0]['width'] filename_bytes = filename.encode('utf-8') image_path = os.path.join( imageDir_path, filename) with open(image_path, 'rb') as file: encoded_jpg = file.read() image_format = b'jpg' xmin_list = list(group['xmin'] / width ) xmax_list = list(group['xmax'] / width ) ymin_list = list(group['ymin'] / height ) ymax_list = list(group['ymax'] / height ) classText_list = [ classText.encode('utf-8') for classText in group['class']] classLabel_list = [ classText_to_classLabel(classText) for classText in group['class']] tf_example = tf.train.Example( features=tf.train.Features( feature = { 'image/height' : dataset_util.int64_feature(height), 'image/width' : dataset_util.int64_feature(width), 'image/filename' : dataset_util.bytes_feature(filename_bytes), 'image/source_id' : dataset_util.bytes_feature(filename_bytes), 'image/encoded' : dataset_util.bytes_feature(encoded_jpg), 'image/format' : dataset_util.bytes_feature(image_format), 'image/object/bbox/xmin' : dataset_util.float_list_feature(xmin_list), 'image/object/bbox/xmax' : dataset_util.float_list_feature(xmax_list), 'image/object/bbox/ymin' : dataset_util.float_list_feature(ymin_list), 'image/object/bbox/ymax' : dataset_util.float_list_feature(ymax_list), 'image/object/class/text' : dataset_util.bytes_list_feature(classText_list), 'image/object/class/label' : dataset_util.int64_list_feature(classLabel_list), })) tfrecord_writer.write(tf_example.SerializeToString()) tfrecord_writer.close() print('成功产生tfrecord文件,保存在路径:%s' % tfrecord_path) #如果训练自己的模型,目标检测类别不同,需要修改此处 def classText_to_classLabel(row_label): if row_label == 'fish': return 1 elif row_label == 'human_face': return 2 else: return None dir_name = 'training' if not os.path.isdir(dir_name): os.mkdir(dir_name) csv2tfrecord('train.csv', 'smaller_images', 'training/train.tfrecord') csv2tfrecord('test.csv', 'smaller_images', 'training/test.tfrecord')
成功界面如下:
四、编写一个配置文件
(1)新建一个文件:D:\ai\objDetect3\mubiaojiance\training\my_label_map.pbtxt
同时输入如下内容:
item { name : "fish" id : 1 }
(2)复制D:\ai\objDetect\models\research\object_detection\samples\configs\ssdlite_mobilenet_v2_coco.config
到 D:\ai\objDetect3\mubiaojiance\training\ssdlite_mobilenet_v2_coco.config
然后打开ssdlite_mobilenet_v2_coco.config进行编辑
第9行的num_classes,对于本文来说,此数设置为1
第143行的batch_size,对于本文来说,此数设置为5,读者根据自己的电脑配置,可以调高或者调低
第177行input_path设置成"training/train.tfrecord"
第179行label_map_path设置成"training/my_label_map.pbtxt"
第191行input_path设置成"training/test.tfrecord"
第193行label_map_path设置成"training/my_label_map.pbtxt"
第158、159这2行需要删除
修改配置文件ssdlite_mobilenet_v2_coco.config并保存后,此时文件夹training中有4个文件,如下图所示:
至此,前置工作基本完成,可以开始训练我们的模型了。
点击进入下一篇
python人工智能 利用Tensorflow object detection API 搭建物体识别模型(三)CPU版