├── images ├── frame1438.jpg └── frame2090.jpg ├── README.md ├── xml_to_csv.py ├── protoc command.txt ├── generate_tfrecord.py ├── object_detection_livestream.py ├── object_detection_video.py └── object_detection.py /images/frame1438.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/XiangGuo1992/Screen-Vehicle-Detection-using-Tensorflow-API/HEAD/images/frame1438.jpg -------------------------------------------------------------------------------- /images/frame2090.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/XiangGuo1992/Screen-Vehicle-Detection-using-Tensorflow-API/HEAD/images/frame2090.jpg -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Screen-Vehicle-Detection-using-Tensorflow-API 2 | Screen&Vehicle Detection in eye tracking videos using Tensorflow API. 3 | All the blogs and videos are in Chinese. 4 | 5 | 博客与视频均为中文。 6 | 7 | ## 1.Build up environment 8 | http://blog.csdn.net/dy_guox/article/details/79081499 9 | 10 | ## 2.Train and Test Customed Model 11 | http://blog.csdn.net/dy_guox/article/details/79111949 12 | 13 | ## 3.Transfer trained model into Android Smartphone 14 | https://blog.csdn.net/dy_guox/article/details/80192343 15 | 16 | 17 | ## Videos (In Chinese) 18 | https://www.bilibili.com/video/av21539370/ 19 | 20 | https://www.bilibili.com/video/av22957279/ 21 | -------------------------------------------------------------------------------- /xml_to_csv.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Tue Jan 16 00:52:02 2018 4 | 5 | @author: Xiang Guo 6 | """ 7 | 8 | import os 9 | import glob 10 | import pandas as pd 11 | import xml.etree.ElementTree as ET 12 | 13 | os.chdir('D:\\test\\test_images\\frame2') 14 | path = 'D:\\test\\test_images\\frame2' 15 | 16 | def xml_to_csv(path): 17 | xml_list = [] 18 | for xml_file in glob.glob(path + '/*.xml'): 19 | tree = ET.parse(xml_file) 20 | root = tree.getroot() 21 | for member in root.findall('object'): 22 | value = (root.find('filename').text, 23 | int(root.find('size')[0].text), 24 | int(root.find('size')[1].text), 25 | member[0].text, 26 | int(member[4][0].text), 27 | int(member[4][1].text), 28 | int(member[4][2].text), 29 | int(member[4][3].text) 30 | ) 31 | xml_list.append(value) 32 | column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] 33 | xml_df = pd.DataFrame(xml_list, columns=column_name) 34 | return xml_df 35 | 36 | 37 | def main(): 38 | image_path = path 39 | xml_df = xml_to_csv(image_path) 40 | xml_df.to_csv('tv_vehicle_labels.csv', index=None) 41 | print('Successfully converted xml to csv.') 42 | 43 | 44 | main() -------------------------------------------------------------------------------- /protoc command.txt: -------------------------------------------------------------------------------- 1 | protoc object_detection/protos/anchor_generator.proto --python_out=. 2 | protoc object_detection/protos/argmax_matcher.proto --python_out=. 3 | protoc object_detection/protos/bipartite_matcher.proto --python_out=. 4 | protoc object_detection/protos/box_coder.proto --python_out=. 5 | protoc object_detection/protos/box_predictor.proto --python_out=. 6 | protoc object_detection/protos/eval.proto --python_out=. 7 | protoc object_detection/protos/faster_rcnn.proto --python_out=. 8 | protoc object_detection/protos/faster_rcnn_box_coder.proto --python_out=. 9 | protoc object_detection/protos/grid_anchor_generator.proto --python_out=. 10 | protoc object_detection/protos/hyperparams.proto --python_out=. 11 | protoc object_detection/protos/image_resizer.proto --python_out=. 12 | protoc object_detection/protos/input_reader.proto --python_out=. 13 | protoc object_detection/protos/keypoint_box_coder.proto --python_out=. 14 | protoc object_detection/protos/losses.proto --python_out=. 15 | protoc object_detection/protos/matcher.proto --python_out=. 16 | protoc object_detection/protos/mean_stddev_box_coder.proto --python_out=. 17 | protoc object_detection/protos/model.proto --python_out=. 18 | protoc object_detection/protos/optimizer.proto --python_out=. 19 | protoc object_detection/protos/pipeline.proto --python_out=. 20 | protoc object_detection/protos/post_processing.proto --python_out=. 21 | protoc object_detection/protos/preprocessor.proto --python_out=. 22 | protoc object_detection/protos/region_similarity_calculator.proto --python_out=. 23 | protoc object_detection/protos/square_box_coder.proto --python_out=. 24 | protoc object_detection/protos/ssd.proto --python_out=. 25 | protoc object_detection/protos/ssd_anchor_generator.proto --python_out=. 26 | protoc object_detection/protos/string_int_label_map.proto --python_out=. 27 | protoc object_detection/protos/train.proto --python_out=. 28 | -------------------------------------------------------------------------------- /generate_tfrecord.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Tue Jan 16 01:04:55 2018 4 | 5 | @author: Xiang Guo 6 | """ 7 | 8 | """ 9 | Usage: 10 | # From tensorflow/models/ 11 | # Create train data: 12 | python generate_tfrecord.py --csv_input=data/tv_vehicle_labels.csv --output_path=train.record 13 | # Create test data: 14 | python generate_tfrecord.py --csv_input=data/test_labels.csv --output_path=test.record 15 | """ 16 | 17 | 18 | 19 | import os 20 | import io 21 | import pandas as pd 22 | import tensorflow as tf 23 | 24 | from PIL import Image 25 | from object_detection.utils import dataset_util 26 | from collections import namedtuple, OrderedDict 27 | 28 | os.chdir('D:\\tensorflow-model\\models\\research\\object_detection\\') 29 | 30 | flags = tf.app.flags 31 | flags.DEFINE_string('csv_input', '', 'Path to the CSV input') 32 | flags.DEFINE_string('output_path', '', 'Path to output TFRecord') 33 | FLAGS = flags.FLAGS 34 | 35 | 36 | # TO-DO replace this with label map 37 | def class_text_to_int(row_label): 38 | if row_label == 'tv': 39 | return 1 40 | elif row_label == 'vehicle': 41 | return 2 42 | else: 43 | None 44 | 45 | 46 | def split(df, group): 47 | data = namedtuple('data', ['filename', 'object']) 48 | gb = df.groupby(group) 49 | return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)] 50 | 51 | 52 | def create_tf_example(group, path): 53 | with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid: 54 | encoded_jpg = fid.read() 55 | encoded_jpg_io = io.BytesIO(encoded_jpg) 56 | image = Image.open(encoded_jpg_io) 57 | width, height = image.size 58 | 59 | filename = group.filename.encode('utf8') 60 | image_format = b'jpg' 61 | xmins = [] 62 | xmaxs = [] 63 | ymins = [] 64 | ymaxs = [] 65 | classes_text = [] 66 | classes = [] 67 | 68 | for index, row in group.object.iterrows(): 69 | xmins.append(row['xmin'] / width) 70 | xmaxs.append(row['xmax'] / width) 71 | ymins.append(row['ymin'] / height) 72 | ymaxs.append(row['ymax'] / height) 73 | classes_text.append(row['class'].encode('utf8')) 74 | classes.append(class_text_to_int(row['class'])) 75 | 76 | tf_example = tf.train.Example(features=tf.train.Features(feature={ 77 | 'image/height': dataset_util.int64_feature(height), 78 | 'image/width': dataset_util.int64_feature(width), 79 | 'image/filename': dataset_util.bytes_feature(filename), 80 | 'image/source_id': dataset_util.bytes_feature(filename), 81 | 'image/encoded': dataset_util.bytes_feature(encoded_jpg), 82 | 'image/format': dataset_util.bytes_feature(image_format), 83 | 'image/object/bbox/xmin': dataset_util.float_list_feature(xmins), 84 | 'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs), 85 | 'image/object/bbox/ymin': dataset_util.float_list_feature(ymins), 86 | 'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs), 87 | 'image/object/class/text': dataset_util.bytes_list_feature(classes_text), 88 | 'image/object/class/label': dataset_util.int64_list_feature(classes), 89 | })) 90 | return tf_example 91 | 92 | 93 | def main(_): 94 | writer = tf.python_io.TFRecordWriter(FLAGS.output_path) 95 | path = os.path.join(os.getcwd(), 'images') 96 | examples = pd.read_csv(FLAGS.csv_input) 97 | grouped = split(examples, 'filename') 98 | for group in grouped: 99 | tf_example = create_tf_example(group, path) 100 | writer.write(tf_example.SerializeToString()) 101 | 102 | writer.close() 103 | output_path = os.path.join(os.getcwd(), FLAGS.output_path) 104 | print('Successfully created the TFRecords: {}'.format(output_path)) 105 | 106 | 107 | if __name__ == '__main__': 108 | tf.app.run() -------------------------------------------------------------------------------- /object_detection_livestream.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Thu Jan 11 16:55:43 2018 4 | 5 | @author: Xiang Guo 6 | """ 7 | #Imports 8 | import time 9 | start = time.time() 10 | import numpy as np 11 | import os 12 | import six.moves.urllib as urllib 13 | import sys 14 | import tarfile 15 | import tensorflow as tf 16 | import zipfile 17 | import cv2 18 | 19 | from collections import defaultdict 20 | from io import StringIO 21 | from matplotlib import pyplot as plt 22 | from PIL import Image 23 | 24 | if tf.__version__ < '1.4.0': 25 | raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!') 26 | 27 | os.chdir('D:\\tensorflow-model\\models\\research\\object_detection') 28 | 29 | 30 | #Env setup 31 | # This is needed to display the images. 32 | #%matplotlib inline 33 | 34 | # This is needed since the notebook is stored in the object_detection folder. 35 | sys.path.append("..") 36 | 37 | 38 | 39 | 40 | #Object detection imports 41 | from utils import label_map_util 42 | 43 | from utils import visualization_utils as vis_util 44 | 45 | 46 | 47 | 48 | #Model preparation 49 | # What model to download. 50 | MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17' #[30,21] best 51 | #MODEL_NAME = 'ssd_inception_v2_coco_2017_11_17' #[42,24] 52 | #MODEL_NAME = 'faster_rcnn_inception_v2_coco_2017_11_08' #[58,28] 53 | #MODEL_NAME = 'faster_rcnn_resnet50_coco_2017_11_08' #[89,30] 54 | #MODEL_NAME = 'faster_rcnn_resnet50_lowproposals_coco_2017_11_08' #[64, ] 55 | #MODEL_NAME = 'rfcn_resnet101_coco_2017_11_08' #[106,32] 56 | 57 | MODEL_FILE = MODEL_NAME + '.tar.gz' 58 | DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/' 59 | 60 | # Path to frozen detection graph. This is the actual model that is used for the object detection. 61 | PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb' 62 | 63 | # List of the strings that is used to add correct label for each box. 64 | PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') 65 | 66 | NUM_CLASSES = 90 67 | 68 | 69 | 70 | ''' 71 | #Download Model 72 | 73 | opener = urllib.request.URLopener() 74 | opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE) 75 | tar_file = tarfile.open(MODEL_FILE) 76 | for file in tar_file.getmembers(): 77 | file_name = os.path.basename(file.name) 78 | if 'frozen_inference_graph.pb' in file_name: 79 | tar_file.extract(file, os.getcwd()) 80 | ''' 81 | 82 | 83 | #Load a (frozen) Tensorflow model into memory. 84 | detection_graph = tf.Graph() 85 | with detection_graph.as_default(): 86 | od_graph_def = tf.GraphDef() 87 | with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: 88 | serialized_graph = fid.read() 89 | od_graph_def.ParseFromString(serialized_graph) 90 | tf.import_graph_def(od_graph_def, name='') 91 | 92 | 93 | #Loading label map 94 | label_map = label_map_util.load_labelmap(PATH_TO_LABELS) 95 | categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) 96 | category_index = label_map_util.create_category_index(categories) 97 | 98 | 99 | #Helper code 100 | def load_image_into_numpy_array(image): 101 | (im_width, im_height) = image.size 102 | return np.array(image.getdata()).reshape( 103 | (im_height, im_width, 3)).astype(np.uint8) 104 | 105 | 106 | #Detection 107 | # For the sake of simplicity we will use only 2 images: 108 | # image1.jpg 109 | # image2.jpg 110 | # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS. 111 | PATH_TO_TEST_IMAGES_DIR = 'test_images2' 112 | TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'frame{}.jpg'.format(i)) for i in range(2176, 2179) ] 113 | 114 | # Size, in inches, of the output images. 115 | IMAGE_SIZE = (12, 8) 116 | 117 | output_path = ('D:\\tensorflow-model\\models\\research\\object_detection\\test_output\\') 118 | 119 | 120 | vidcap = cv2.VideoCapture(0) 121 | 122 | with detection_graph.as_default(): 123 | with tf.Session(graph=detection_graph) as sess: 124 | # Definite input and output Tensors for detection_graph 125 | image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') 126 | # Each box represents a part of the image where a particular object was detected. 127 | detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') 128 | # Each score represent how level of confidence for each of the objects. 129 | # Score is shown on the result image, together with the class label. 130 | detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') 131 | detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') 132 | num_detections = detection_graph.get_tensor_by_name('num_detections:0') 133 | 134 | while(True): 135 | ret, image_np = vidcap.read() 136 | 137 | if ret == True: 138 | 139 | 140 | 141 | # Expand dimensions since the model expects images to have shape: [1, None, None, 3] 142 | image_np_expanded = np.expand_dims(image_np, axis=0) 143 | # Actual detection. 144 | (boxes, scores, classes, num) = sess.run( 145 | [detection_boxes, detection_scores, detection_classes, num_detections], 146 | feed_dict={image_tensor: image_np_expanded}) 147 | # Visualization of the results of a detection. 148 | vis_util.visualize_boxes_and_labels_on_image_array( 149 | image_np, 150 | np.squeeze(boxes), 151 | np.squeeze(classes).astype(np.int32), 152 | np.squeeze(scores), 153 | category_index, 154 | use_normalized_coordinates=True, 155 | line_thickness=8) 156 | 157 | cv2.imshow('object_detection',cv2.resize(image_np,(800,600))) 158 | if cv2.waitKey(25) & 0xFF == ord('q'): 159 | cv2.destroyAllWindows() 160 | break 161 | 162 | 163 | 164 | 165 | # Break the loop 166 | else: 167 | break 168 | 169 | 170 | end = time.time() 171 | print("Execution Time: ", end - start) 172 | 173 | 174 | -------------------------------------------------------------------------------- /object_detection_video.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Thu Jan 11 16:55:43 2018 4 | 5 | @author: Xiang Guo 6 | """ 7 | #Imports 8 | import time 9 | start = time.time() 10 | import numpy as np 11 | import os 12 | import six.moves.urllib as urllib 13 | import sys 14 | import tarfile 15 | import tensorflow as tf 16 | import zipfile 17 | import cv2 18 | 19 | from collections import defaultdict 20 | from io import StringIO 21 | from matplotlib import pyplot as plt 22 | from PIL import Image 23 | 24 | if tf.__version__ < '1.4.0': 25 | raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!') 26 | 27 | os.chdir('D:\\tensorflow-model\\models\\research\\object_detection') 28 | 29 | 30 | #Env setup 31 | # This is needed to display the images. 32 | #%matplotlib inline 33 | 34 | # This is needed since the notebook is stored in the object_detection folder. 35 | sys.path.append("..") 36 | 37 | 38 | 39 | 40 | #Object detection imports 41 | from utils import label_map_util 42 | 43 | from utils import visualization_utils as vis_util 44 | 45 | 46 | 47 | 48 | #Model preparation 49 | # What model to download. 50 | MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17' #[30,21] best 51 | #MODEL_NAME = 'ssd_inception_v2_coco_2017_11_17' #[42,24] 52 | #MODEL_NAME = 'faster_rcnn_inception_v2_coco_2017_11_08' #[58,28] 53 | #MODEL_NAME = 'faster_rcnn_resnet50_coco_2017_11_08' #[89,30] 54 | #MODEL_NAME = 'faster_rcnn_resnet50_lowproposals_coco_2017_11_08' #[64, ] 55 | #MODEL_NAME = 'rfcn_resnet101_coco_2017_11_08' #[106,32] 56 | 57 | MODEL_FILE = MODEL_NAME + '.tar.gz' 58 | DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/' 59 | 60 | # Path to frozen detection graph. This is the actual model that is used for the object detection. 61 | PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb' 62 | 63 | # List of the strings that is used to add correct label for each box. 64 | PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') 65 | 66 | NUM_CLASSES = 90 67 | 68 | 69 | 70 | ''' 71 | #Download Model 72 | 73 | opener = urllib.request.URLopener() 74 | opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE) 75 | tar_file = tarfile.open(MODEL_FILE) 76 | for file in tar_file.getmembers(): 77 | file_name = os.path.basename(file.name) 78 | if 'frozen_inference_graph.pb' in file_name: 79 | tar_file.extract(file, os.getcwd()) 80 | ''' 81 | 82 | 83 | #Load a (frozen) Tensorflow model into memory. 84 | detection_graph = tf.Graph() 85 | with detection_graph.as_default(): 86 | od_graph_def = tf.GraphDef() 87 | with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: 88 | serialized_graph = fid.read() 89 | od_graph_def.ParseFromString(serialized_graph) 90 | tf.import_graph_def(od_graph_def, name='') 91 | 92 | 93 | #Loading label map 94 | label_map = label_map_util.load_labelmap(PATH_TO_LABELS) 95 | categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) 96 | category_index = label_map_util.create_category_index(categories) 97 | 98 | 99 | #Helper code 100 | def load_image_into_numpy_array(image): 101 | (im_width, im_height) = image.size 102 | return np.array(image.getdata()).reshape( 103 | (im_height, im_width, 3)).astype(np.uint8) 104 | 105 | 106 | #Detection 107 | # For the sake of simplicity we will use only 2 images: 108 | # image1.jpg 109 | # image2.jpg 110 | # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS. 111 | PATH_TO_TEST_IMAGES_DIR = 'test_images2' 112 | TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'frame{}.jpg'.format(i)) for i in range(2176, 2179) ] 113 | 114 | # Size, in inches, of the output images. 115 | IMAGE_SIZE = (12, 8) 116 | 117 | output_path = ('D:\\tensorflow-model\\models\\research\\object_detection\\test_output\\') 118 | 119 | 120 | 121 | 122 | with detection_graph.as_default(): 123 | with tf.Session(graph=detection_graph) as sess: 124 | # Definite input and output Tensors for detection_graph 125 | image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') 126 | # Each box represents a part of the image where a particular object was detected. 127 | detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') 128 | # Each score represent how level of confidence for each of the objects. 129 | # Score is shown on the result image, together with the class label. 130 | detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') 131 | detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') 132 | num_detections = detection_graph.get_tensor_by_name('num_detections:0') 133 | 134 | #the video to be detected, eg, "test.mp4" here 135 | vidcap = cv2.VideoCapture('D:\\tensorflow-model\\models\\research\\object_detection\\test.mp4') 136 | # Default resolutions of the frame are obtained.The default resolutions are system dependent. 137 | # We convert the resolutions from float to integer. 138 | frame_width = int(vidcap.get(3)) 139 | frame_height = int(vidcap.get(4)) 140 | 141 | # Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file. 142 | out_video = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 10, (frame_width,frame_height)) 143 | 144 | while(True): 145 | ret, image = vidcap.read() 146 | 147 | if ret == True: 148 | 149 | #image_np = load_image_into_numpy_array(image) 150 | image_np = image 151 | 152 | # Expand dimensions since the model expects images to have shape: [1, None, None, 3] 153 | image_np_expanded = np.expand_dims(image_np, axis=0) 154 | # Actual detection. 155 | (boxes, scores, classes, num) = sess.run( 156 | [detection_boxes, detection_scores, detection_classes, num_detections], 157 | feed_dict={image_tensor: image_np_expanded}) 158 | # Visualization of the results of a detection. 159 | vis_util.visualize_boxes_and_labels_on_image_array( 160 | image_np, 161 | np.squeeze(boxes), 162 | np.squeeze(classes).astype(np.int32), 163 | np.squeeze(scores), 164 | category_index, 165 | use_normalized_coordinates=True, 166 | line_thickness=8) 167 | plt.figure(figsize=IMAGE_SIZE) 168 | plt.imshow(image_np) 169 | # Write the frame into the file 'output.avi' 170 | out_video.write(image_np) 171 | 172 | 173 | 174 | # Break the loop 175 | else: 176 | break 177 | 178 | end = time.time() 179 | print("Execution Time: ", end - start) 180 | 181 | 182 | -------------------------------------------------------------------------------- /object_detection.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Thu Jan 11 16:55:43 2018 4 | 5 | @author: Xiang Guo 6 | """ 7 | #Imports 8 | import time 9 | start = time.time() 10 | import numpy as np 11 | import os 12 | import six.moves.urllib as urllib 13 | import sys 14 | import tarfile 15 | import tensorflow as tf 16 | import zipfile 17 | import cv2 18 | 19 | from collections import defaultdict 20 | from io import StringIO 21 | from matplotlib import pyplot as plt 22 | from PIL import Image 23 | import pandas as pd 24 | 25 | if tf.__version__ < '1.4.0': 26 | raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!') 27 | 28 | os.chdir('E:\\tensorflow_models\\models\\research\\object_detection\\') 29 | 30 | 31 | #Env setup 32 | # This is needed to display the images. 33 | #%matplotlib inline 34 | 35 | # This is needed since the notebook is stored in the object_detection folder. 36 | sys.path.append("..") 37 | 38 | 39 | 40 | 41 | #Object detection imports 42 | from object_detection.utils import label_map_util 43 | 44 | 45 | from object_detection.utils import visualization_utils as vis_util 46 | 47 | 48 | 49 | #Model preparation 50 | # What model to download. 51 | #MODEL_NAME = 'tv_vehicle_inference_graph' 52 | #MODEL_NAME = 'tv_vehicle_inference_graph_fasterCNN' 53 | MODEL_NAME = 'tv_vehicle_inference_graph_ssd_mobile' 54 | #MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17' #[30,21] best 55 | #MODEL_NAME = 'ssd_inception_v2_coco_2017_11_17' #[42,24] 56 | #MODEL_NAME = 'faster_rcnn_inception_v2_coco_2017_11_08' #[58,28] 57 | #MODEL_NAME = 'faster_rcnn_resnet50_coco_2017_11_08' #[89,30] 58 | #MODEL_NAME = 'faster_rcnn_resnet50_lowproposals_coco_2017_11_08' #[64, ] 59 | #MODEL_NAME = 'rfcn_resnet101_coco_2017_11_08' #[106,32] 60 | 61 | ''' 62 | MODEL_FILE = MODEL_NAME + '.tar.gz' 63 | DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/' 64 | ''' 65 | 66 | # Path to frozen detection graph. This is the actual model that is used for the object detection. 67 | PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb' 68 | 69 | # List of the strings that is used to add correct label for each box. 70 | PATH_TO_LABELS = os.path.join('training', 'tv_vehicle_detection.pbtxt') 71 | 72 | NUM_CLASSES = 2 73 | 74 | 75 | 76 | ''' 77 | #Download Model 78 | 79 | opener = urllib.request.URLopener() 80 | opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE) 81 | tar_file = tarfile.open(MODEL_FILE) 82 | for file in tar_file.getmembers(): 83 | file_name = os.path.basename(file.name) 84 | if 'frozen_inference_graph.pb' in file_name: 85 | tar_file.extract(file, os.getcwd()) 86 | ''' 87 | 88 | 89 | #Load a (frozen) Tensorflow model into memory. 90 | detection_graph = tf.Graph() 91 | with detection_graph.as_default(): 92 | od_graph_def = tf.GraphDef() 93 | with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: 94 | serialized_graph = fid.read() 95 | od_graph_def.ParseFromString(serialized_graph) 96 | tf.import_graph_def(od_graph_def, name='') 97 | 98 | 99 | #Loading label map 100 | label_map = label_map_util.load_labelmap(PATH_TO_LABELS) 101 | categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) 102 | category_index = label_map_util.create_category_index(categories) 103 | 104 | 105 | #Helper code 106 | def load_image_into_numpy_array(image): 107 | (im_width, im_height) = image.size 108 | return np.array(image.getdata()).reshape( 109 | (im_height, im_width, 3)).astype(np.uint8) 110 | 111 | 112 | #Detection 113 | # For the sake of simplicity we will use only 2 images: 114 | # image1.jpg 115 | # image2.jpg 116 | # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS. 117 | PATH_TO_TEST_IMAGES_DIR = 'G:\\Resource for Xiang\\Lian Cui experiment\\eyetracking data\\analysis\\framepics\\' 118 | os.chdir(PATH_TO_TEST_IMAGES_DIR) 119 | TEST_IMAGE_DIRS = os.listdir(PATH_TO_TEST_IMAGES_DIR) 120 | 121 | # Size, in inches, of the output images. 122 | IMAGE_SIZE = (12, 8) 123 | 124 | output_image_path = ('G:\\Resource for Xiang\\Lian Cui experiment\\eyetracking data\\analysis\\output\\pics\\') 125 | output_csv_path = ('G:\\Resource for Xiang\\Lian Cui experiment\\eyetracking data\\analysis\\output\\csv\\') 126 | 127 | for image_folder in TEST_IMAGE_DIRS: 128 | with detection_graph.as_default(): 129 | with tf.Session(graph=detection_graph) as sess: 130 | # Definite input and output Tensors for detection_graph 131 | image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') 132 | # Each box represents a part of the image where a particular object was detected. 133 | detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') 134 | # Each score represent how level of confidence for each of the objects. 135 | # Score is shown on the result image, together with the class label. 136 | detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') 137 | detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') 138 | num_detections = detection_graph.get_tensor_by_name('num_detections:0') 139 | TEST_IMAGE_PATHS = os.listdir(os.path.join(image_folder)) 140 | os.makedirs(output_image_path+image_folder) 141 | data = pd.DataFrame() 142 | for image_path in TEST_IMAGE_PATHS: 143 | image = Image.open(image_folder + '//'+image_path) 144 | width, height = image.size 145 | # the array based representation of the image will be used later in order to prepare the 146 | # result image with boxes and labels on it. 147 | image_np = load_image_into_numpy_array(image) 148 | # Expand dimensions since the model expects images to have shape: [1, None, None, 3] 149 | image_np_expanded = np.expand_dims(image_np, axis=0) 150 | # Actual detection. 151 | (boxes, scores, classes, num) = sess.run( 152 | [detection_boxes, detection_scores, detection_classes, num_detections], 153 | feed_dict={image_tensor: image_np_expanded}) 154 | # Visualization of the results of a detection. 155 | vis_util.visualize_boxes_and_labels_on_image_array( 156 | image_np, 157 | np.squeeze(boxes), 158 | np.squeeze(classes).astype(np.int32), 159 | np.squeeze(scores), 160 | category_index, 161 | use_normalized_coordinates=True, 162 | line_thickness=8) 163 | #write images 164 | #保存识别结果图片 165 | 166 | cv2.imwrite(output_image_path+image_folder+'\\'+image_path.split('\\')[-1],image_np) 167 | 168 | s_boxes = boxes[scores > 0.5] 169 | s_classes = classes[scores > 0.5] 170 | s_scores=scores[scores>0.5] 171 | 172 | #write table 173 | #保存位置坐标结果到 .csv表格 174 | for i in range(len(s_classes)): 175 | 176 | newdata= pd.DataFrame(0, index=range(1), columns=range(7)) 177 | newdata.iloc[0,0] = image_path.split("\\")[-1].split('.')[0] 178 | newdata.iloc[0,1] = s_boxes[i][0]*height #ymin 179 | newdata.iloc[0,2] = s_boxes[i][1]*width #xmin 180 | newdata.iloc[0,3] = s_boxes[i][2]*height #ymax 181 | newdata.iloc[0,4] = s_boxes[i][3]*width #xmax 182 | newdata.iloc[0,5] = s_scores[i] 183 | newdata.iloc[0,6] = s_classes[i] 184 | 185 | data = data.append(newdata) 186 | data.to_csv(output_csv_path+image_folder+'.csv',index = False) 187 | 188 | end = time.time() 189 | print("Execution Time: ", end - start) 190 | 191 | 192 | --------------------------------------------------------------------------------