Spaces:
Runtime error
Runtime error
| """ | |
| DeepLabCut Toolbox (deeplabcut.org) | |
| © A. & M. Mathis Labs | |
| Licensed under GNU Lesser General Public License v3.0 | |
| """ | |
| import tensorflow as tf | |
| vers = (tf.__version__).split(".") | |
| if int(vers[0]) == 2 or int(vers[0]) == 1 and int(vers[1]) > 12: | |
| tf = tf.compat.v1 | |
| else: | |
| tf = tf | |
| def read_graph(file): | |
| """ | |
| Loads the graph from a protobuf file | |
| Parameters | |
| ----------- | |
| file : string | |
| path to the protobuf file | |
| Returns | |
| -------- | |
| graph_def :class:`tensorflow.tf.compat.v1.GraphDef` | |
| The graph definition of the DeepLabCut model found at the object's path | |
| """ | |
| with tf.io.gfile.GFile(file, "rb") as f: | |
| graph_def = tf.GraphDef() | |
| graph_def.ParseFromString(f.read()) | |
| return graph_def | |
| def finalize_graph(graph_def): | |
| """ | |
| Finalize the graph and get inputs to model | |
| Parameters | |
| ----------- | |
| graph_def :class:`tensorflow.compat.v1.GraphDef` | |
| The graph of the DeepLabCut model, read using the :func:`read_graph` method | |
| Returns | |
| -------- | |
| graph :class:`tensorflow.compat.v1.GraphDef` | |
| The finalized graph of the DeepLabCut model | |
| inputs :class:`tensorflow.Tensor` | |
| Input tensor(s) for the model | |
| """ | |
| graph = tf.Graph() | |
| with graph.as_default(): | |
| tf.import_graph_def(graph_def, name="DLC") | |
| graph.finalize() | |
| return graph | |
| def get_output_nodes(graph): | |
| """ | |
| Get the output node names from a graph | |
| Parameters | |
| ----------- | |
| graph :class:`tensorflow.Graph` | |
| The graph of the DeepLabCut model | |
| Returns | |
| -------- | |
| output : list | |
| the output node names as a list of strings | |
| """ | |
| op_names = [str(op.name) for op in graph.get_operations()] | |
| if "concat_1" in op_names[-1]: | |
| output = [op_names[-1]] | |
| else: | |
| output = [op_names[-1], op_names[-2]] | |
| return output | |
| def get_output_tensors(graph): | |
| """ | |
| Get the names of the output tensors from a graph | |
| Parameters | |
| ----------- | |
| graph :class:`tensorflow.Graph` | |
| The graph of the DeepLabCut model | |
| Returns | |
| -------- | |
| output : list | |
| the output tensor names as a list of strings | |
| """ | |
| output_nodes = get_output_nodes(graph) | |
| output_tensor = [out + ":0" for out in output_nodes] | |
| return output_tensor | |
| def get_input_tensor(graph): | |
| input_tensor = str(graph.get_operations()[0].name) + ":0" | |
| return input_tensor | |
| def extract_graph(graph, tf_config=None): | |
| """ | |
| Initializes a tensorflow session with the specified graph and extracts the model's inputs and outputs | |
| Parameters | |
| ----------- | |
| graph :class:`tensorflow.Graph` | |
| a tensorflow graph containing the desired model | |
| tf_config :class:`tensorflow.ConfigProto` | |
| Returns | |
| -------- | |
| sess :class:`tensorflow.Session` | |
| a tensorflow session with the specified graph definition | |
| outputs :class:`tensorflow.Tensor` | |
| the output tensor(s) for the model | |
| """ | |
| input_tensor = get_input_tensor(graph) | |
| output_tensor = get_output_tensors(graph) | |
| sess = tf.Session(graph=graph, config=tf_config) | |
| inputs = graph.get_tensor_by_name(input_tensor) | |
| outputs = [graph.get_tensor_by_name(out) for out in output_tensor] | |
| return sess, inputs, outputs | |