def __init__(self, imagePath, modelPath):
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
# Name of the directory containing the object detection module we're using
self.MODEL_NAME = modelPath
self.IMAGE_NAME = imagePath
#print(self.IMAGE_NAME)
# Grab path to current working directory
CWD_PATH = os.getcwd()
# Path to frozen detection graph .pb file, which contains the model that is used
# for object detection.
self.PATH_TO_CKPT = os.path.join(CWD_PATH, self.MODEL_NAME, 'frozen_inference_graph.pb')
# Path to label map file
self.PATH_TO_LABELS = os.path.join(CWD_PATH, 'research/data', 'labelmap.pbtxt')
#self.PATH_TO_LABELS = "data/labelmap.pbtxt"
# Path to images
self.PATH_TO_IMAGE = os.path.join(CWD_PATH, 'research', self.IMAGE_NAME)
print(self.PATH_TO_IMAGE)
# Number of classes the object detector can identify
self.NUM_CLASSES = 11
# Load the label map.
# Label maps map indices to category names, so that when our convolution
# network predicts `5`, we know that this corresponds to `king`.
# Here we use internal utility functions, but anything that returns a
# dictionary mapping integers to appropriate string labels would be fine
self.label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
self.categories = label_map_util.convert_label_map_to_categories(self.label_map,
max_num_classes=self.NUM_CLASSES,
use_display_name=True)
self.category_index = label_map_util.create_category_index(self.categories)
self.class_names_mapping = {
1: "ear_ring", 2: "GirlsTopWear", 3: "glass", 4: "hat", 5: "Jacket", 6: "MensShorts", 7: "MensTopWear",
8: "Pant", 9: "shoes", 10: "tie", 11: "watch"
}
self.detection_graph = tf.Graph()
with self.detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
# Define input and output tensors (i.e. data) for the object detection classifier
# Input tensor is the image
self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
# Output tensors are the detection boxes, scores, and classes
# Each box represents a part of the image where a particular object was detected
self.detection_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represents level of confidence for each of the objects.
# The score is shown on the result image, together with the class label.
self.detection_scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
self.detection_classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
# Number of objects detected
self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
def test_load_bad_label_map(self):
label_map_string = """
item {
id:0
name:'class that should not be indexed at zero'
}
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
with self.assertRaises(ValueError):
label_map_util.load_labelmap(label_map_path)
def __init__(self, imagename, modelpath):
self.IMAGE = imagename
self.MODEL_NAME = modelpath
sys.path.append("..")
CWD_path = os.getcwd()
self.PATH_TO_CKPT = os.path.join(CWD_path, self.MODEL_NAME,
'frozen_inference_graph.pb')
self.PATH_TO_LABELS = os.path.join(CWD_path, 'research/data',
'labelmap.pbtxt')
self.PATH_TO_IMAGE = os.path.join(CWD_path, 'research', self.IMAGE)
self.NUM_OF_CLASSES = 1
self.label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
self.categories = label_map_util.convert_label_map_to_categories(
self.label_map,
max_num_classes=self.NUM_OF_CLASSES,
use_display_name=True)
self.category_index = label_map_util.create_category_index(
self.categories)
self.class_names_mapping = {1: "Helmets"}
self.detection_graph = tf.Graph()
with self.detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
# Input tensor is the image
self.image_tensor = self.detection_graph.get_tensor_by_name(
'image_tensor:0')
# Output tensors are the detection boxes, scores, and classes
# Each box represents a part of the image where a particular object was detected
self.detection_boxes = self.detection_graph.get_tensor_by_name(
'detection_boxes:0')
# Each score represents level of confidence for each of the objects.
# The score is shown on the result image, together with the class label.
self.detection_scores = self.detection_graph.get_tensor_by_name(
'detection_scores:0')
self.detection_classes = self.detection_graph.get_tensor_by_name(
'detection_classes:0')
# Number of objects detected
self.num_detections = self.detection_graph.get_tensor_by_name(
'num_detections:0')
def __init__(self, imagePath, modelPath):
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
# Name of the directory containing the object detection module we're using
self.MODEL_NAME = modelPath
self.IMAGE_NAME = imagePath
# print(self.IMAGE_NAME)
# Grab path to current working directory
CWD_PATH = os.getcwd()
# Path to frozen detection graph .pb file, which contains the model that is used
# for object detection.
self.PATH_TO_CKPT = os.path.join(CWD_PATH, self.MODEL_NAME, 'frozen_inference_graph.pb')
# Path to label map file
self.PATH_TO_LABELS = os.path.join(CWD_PATH, 'research/data', 'labelmap.pbtxt')
# self.PATH_TO_LABELS = "data/labelmap.pbtxt"
# Path to images
self.PATH_TO_IMAGE = os.path.join(CWD_PATH, 'research', self.IMAGE_NAME)
print(self.PATH_TO_IMAGE)
# Number of classes the object detector can identify
self.NUM_CLASSES = 1
# Load the label map.
# Label maps map indices to category names, so that when our convolution
# network predicts `5`, we know that this corresponds to `king`.
# Here we use internal utility functions, but anything that returns a
# dictionary mapping integers to appropriate string labels would be fine
self.label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
self.categories = label_map_util.convert_label_map_to_categories(self.label_map,
max_num_classes=self.NUM_CLASSES,
use_display_name=True)
self.category_index = label_map_util.create_category_index(self.categories)
'''
self.class_names_mapping = {
1: "person", 2: "bicycle", 3: "car", 4: "motorcycle", 5: "airplane", 6: "bus", 7: "train", 8: "truck",
9: "boat",
10: "traffic light",
11: "fire hydrant", 13: "stop sign", 14: "parking meter", 15: "bench", 16: "bird", 17: "cat", 18: "dog",
19: "horse", 20: "sheep",
21: "cow", 22: "elephant", 23: "bear", 24: "zebra", 25: "giraffe", 27: "backpack", 28: "umbrella",
31: "handbag",
32: "tie", 33: "suitcase",
34: "frisbee", 35: "skis", 36: "snowboard", 37: "sports ball", 38: "kite", 39: "baseball bat",
40: "baseball glove",
41: "skateboard", 42: "surfboard", 43: "tennis racket", 44: "bottle",
46: "wine glass", 47: "cup", 48: "fork", 49: "knife", 50: "spoon", 51: "bowl", 52: "banana", 53: "apple",
54: "sandwich", 55: "orange", 56: "broccoli", 57: "carrot", 58: "hot dog", 59: "pizza", 60: "donut",
61: "cake", 62: "chair", 63: "couch", 64: "potted plant", 65: "bed", 67: "dining table", 70: "toilet",
72: "tv",
73: "laptop", 74: "mouse", 75: "remote", 76: "keyboard", 77: "cell", 78: "microwave", 79: "oven",
80: "toaster",
81: "sink", 82: "refrigerator", 84: "book", 85: "clock", 86: "vase", 87: "scissors", 88: "teddy bear",
89: "hair drier", 90: "toothbrush"
}'''
self.class_names_mapping = {1: "Helmets"}
self.detection_graph = tf.Graph()
with self.detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
# Define input and output tensors (i.e. data) for the object detection classifier
# Input tensor is the image
self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
# Output tensors are the detection boxes, scores, and classes
# Each box represents a part of the image where a particular object was detected
self.detection_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represents level of confidence for each of the objects.
# The score is shown on the result image, together with the class label.
self.detection_scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
self.detection_classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
# Number of objects detected
self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
file_name = os.path.basename(file.name)
if 'frozen_inference_graph.pb' in file_name:
tar_file.extract(file, os.getcwd())
"""
# Load a (frozen) Tensorflow model into memory.
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='')
# Loading label 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)
# Helper code
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)
# For the sake of simplicity we will use only 2 images:
# image1.jpg
# image2.jpg
def _load_label_map(self):
label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(
label_map, max_num_classes=self.NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
return category_index
