Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import numpy as np
- import os
- import six.moves.urllib as urllib
- import sys
- import tarfile
- import tensorflow as tf
- import zipfile
- from collections import defaultdict
- from io import StringIO
- from matplotlib import pyplot as plt
- from PIL import Image
- from object_detection.utils import label_map_util
- from object_detection.utils import visualization_utils as vis_util
- PATH_TO_CKPT = 'frozen_graph/frozen_inference_graph.pb'
- PATH_TO_LABELS = 'object_detection/pascal_label_map.pbtxt'
- NUM_CLASSES = 7
- 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='')
- 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)
- def load_image_into_numpy_array(image):
- (im_width, im_height) = image.size
- return np.array(image.getdata()).reshape(
- (im_height, im_width, 3)).astype(np.uint8)
- PATH_TO_TEST_IMAGES_DIR = 'object_detection/test_images/'
- TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 2) ]
- IMAGE_SIZE = (12, 8)
- with detection_graph.as_default():
- with tf.Session(graph=detection_graph) as sess:
- sess.run(tf.global_variables_initializer())
- img = 1
- for image_path in TEST_IMAGE_PATHS:
- image = Image.open(image_path)
- image_np = load_image_into_numpy_array(image)
- # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
- image_np_expanded = np.expand_dims(image_np, axis=0)
- 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.
- boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
- scores = detection_graph.get_tensor_by_name('detection_scores:0')
- classes = detection_graph.get_tensor_by_name('detection_classes:0')
- num_detections = detection_graph.get_tensor_by_name('num_detections:0')
- (boxes, scores, classes, num_detections) = sess.run(
- [boxes, scores, classes, num_detections],
- feed_dict={image_tensor: image_np_expanded})
- 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)
- plt.figure(figsize=IMAGE_SIZE)
- plt.imsave('RESULTS/' + str(img) + '.jpg', image_np)
- img += 1
- # Return found objects
- print([category_index.get(i) for i in classes[0]])
- print(boxes.shape)
- print(num_detections)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement