def __init__(self, name: str, label_path: pathlib.Path):
self.name: str = name
label_map = label_map_util.load_labelmap(str(label_path))
self.categories = label_map_util.convert_label_map_to_categories(
label_map, max_num_classes=1000)
self.category_index: CategoryIndex = (
label_map_util.create_category_index(self.categories))
def load(cls, graph_path: pathlib.Path, label_path: pathlib.Path):
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(str(graph_path), "rb") as f:
serialized_graph = f.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name="")
label_map = label_map_util.load_labelmap(str(label_path))
logger.info("Model loaded")
return cls(graph=detection_graph, label_map=label_map)