seresnet50训练自己的数据集_训练自己的数据集

#-*- coding: utf-8 -*-

对tutorial做了一点修改，改成py文件importtime

start=time.time()importnumpy as npimportosimportsix.moves.urllib as urllibimportsysimporttarfileimporttensorflow as tfimportzipfileimportcv2from collections importdefaultdictfrom io importStringIOfrom matplotlib importpyplot as pltfrom PIL importImageimportpandas as pdif tf.__version__ < '1.4.0':raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')#改变当前工作目录到指定的路径#默认情况下，Python解释器会搜索当前目录、所有已安装的内置模块和第三方模块#sys.path,返回的模块搜索路径列表，类似于这样的：#[#'F:\\', 'F:\\测试自己的图像识别模型.py',#'D:\\anaconda\\anaconda3.4.2.0\\python35.zip',#'D:\\anaconda\\anaconda3.4.2.0\\DLLs',#'D:\\anaconda\\anaconda3.4.2.0\\lib'#]

os.chdir(r'E:/object_detection/')#.. 或 ../是父目录，表示将父目录添加进了sys.path#. 或 ./是当前目录

sys.path.append("..")from object_detection.utils importlabel_map_utilfrom object_detection.utils importvisualization_utils as vis_util#Model preparation#What model to download.#MODEL_NAME = 'tv_vehicle_inference_graph'#MODEL_NAME = 'tv_vehicle_inference_graph_fasterCNN'

#最后生成的那个文件夹名result

MODEL_NAME = 'result'

#MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17' #[30,21] best#MODEL_NAME = 'ssd_inception_v2_coco_2017_11_17' #[42,24]#MODEL_NAME = 'faster_rcnn_inception_v2_coco_2017_11_08' #[58,28]#MODEL_NAME = 'faster_rcnn_resnet50_coco_2017_11_08' #[89,30]#MODEL_NAME = 'faster_rcnn_resnet50_lowproposals_coco_2017_11_08' #[64, ]#MODEL_NAME = 'rfcn_resnet101_coco_2017_11_08' #[106,32]

#PATH_TO_CKPT = 'result' + '/你自己生成的pb模型文件名'

PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

#标签问价你在training文件夹下面，当前脚本和training文件夹是同一级的

PATH_TO_LABELS = os.path.join('training', 'detection.pbtxt')#目标类倍数

NUM_CLASSES = 2

#将训练完的载入内存(不用改)

detection_graph =tf.Graph()

with detection_graph.as_default():

od_graph_def=tf.GraphDef()

with tf.gfile.GFile(PATH_TO_CKPT,'rb') as fid:

serialized_graph=fid.read()

od_graph_def.ParseFromString(serialized_graph)

tf.import_graph_def(od_graph_def, name='')#载入标签map(不用改)

label_map =label_map_util.load_labelmap(PATH_TO_LABELS)

categories= label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)

category_index=label_map_util.create_category_index(categories)defload_image_into_numpy_array(image):

(im_width, im_height)=image.size#getdata()返回的是2维度的：(width*height , 3)

#[[1 2 3]

#[3 4 5]]类似这种的

#因此reshape成3维

return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)#放图片文件夹的的文件夹名!!!就是test_images文件夹在PATH_TO_TEST_IMAGES_DIR文件夹里面#这个脚本和 PATH_TO_TEST_IMAGES_DIR 文件夹 是在同一个目录下的，都在object_detection/下面

PATH_TO_TEST_IMAGES_DIR = 'PATH_TO_TEST_IMAGES_DIR'

#将PATH_TO_TEST_IMAGES_DIR文件夹目录设置为当前路径

os.chdir(PATH_TO_TEST_IMAGES_DIR)#os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表，带后缀#这个列表以字母顺序。 它不包括 '.' 和'..' 即使它在文件夹中。#这个列表里就只有一个元素：test_images文件夹

TEST_IMAGE_DIRS =os.listdir(PATH_TO_TEST_IMAGES_DIR)#输出图片的尺寸，inches

IMAGE_SIZE = (12, 8)#输出图片带画框的结果，这个文件夹和test_images同级

output_image_path = ("PATH_TO_TEST_IMAGES_DIR\")

#另外加了输出识别结果框的坐标，保存为.csv表格文件

output_csv_path = ("\输出\识别结果\表格\的\路径\")

#image_folder：test_images文件夹

for image_folder inTEST_IMAGE_DIRS:

with detection_graph.as_default():

with tf.Session(graph=detection_graph) as sess:#Definite input and output Tensors for detection_graph

image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')#Each box represents a part of the image where a particular object was detected.

detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')#Each score represent how level of confidence for each of the objects.

#Score is shown on the result image, together with the class label.

detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')

detection_classes= detection_graph.get_tensor_by_name('detection_classes:0')

num_detections= detection_graph.get_tensor_by_name('num_detections:0')#当前路径加test_images文件夹名，再listdir得到的就是所有test图片的path列表。

TEST_IMAGE_PATHS =os.listdir(os.path.join(image_folder))#新建一个文件夹PATH_TO_TEST_IMAGES_DIR\draw_test_images

os.makedirs(output_image_path+'draw_'+image_folder)

data=pd.DataFrame()for image_path inTEST_IMAGE_PATHS:

image= Image.open(image_folder + '//'+image_path)

width, height=image.size#the array based representation of the image will be used later in order to prepare the

#result image with boxes and labels on it.

image_np = load_image_into_numpy_array(image)#[width,height,2]

#在axis=0即第一个维度扩充一个维度，就是在最外面加了一层括号,shpe=[1,width,height,3]

#这样变是因为，输入的图片格式要是四维的

image_np_expanded = np.expand_dims(image_np, axis=0)#run检测，得到结果值

(boxes, scores, classes, num) =sess.run([detection_boxes, detection_scores, detection_classes, num_detections],

feed_dict={image_tensor: image_np_expanded})#可视化

#vis_util是objection模块里的一个函数

vis_util.visualize_boxes_and_labels_on_image_array(

image_np,

np.squeeze(boxes),

np.squeeze(classes).astype(np.int32),

np.squeeze(scores),

category_index,

use_normalized_coordinates=True,

line_thickness=8)#保存识别结果图片，输出的也是图片，保存在draw_test_images文件夹下

cv2.imwrite(output_image_path+'draw_'+image_folder+'\\'+image_path.split('\\')[-1],image_np)

s_boxes= boxes[scores > 0.5]

s_classes= classes[scores > 0.5]

s_scores=scores[scores>0.5]#write table

#保存位置坐标结果到 .csv表格

for i inrange(len(s_classes)):

newdata= pd.DataFrame(0, index=range(1), columns=range(7))

newdata.iloc[0,0]= image_path.split("\\")[-1].split('.')[0]

newdata.iloc[0,1] = s_boxes[i][0]*height #ymin

newdata.iloc[0,2] = s_boxes[i][1]*width #xmin

newdata.iloc[0,3] = s_boxes[i][2]*height #ymax

newdata.iloc[0,4] = s_boxes[i][3]*width #xmax

newdata.iloc[0,5] =s_scores[i]

newdata.iloc[0,6] =s_classes[i]

data=data.append(newdata)

data.to_csv(output_csv_path+image_folder+'.csv',index =False)

end=time.time()print("Execution Time:", end - start)