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)