metadata parsing

onnx/src/model.rs

@@ -1,5 +1,5 @@
 use std::path::PathBuf;
-use std::{fs, path};
+use std::{ fs, path};
 
 use std::collections::HashMap;
 
@@ -8,7 +8,7 @@ use tract_hir::prelude::tract_itertools::Itertools;
 
 use crate::data_resolver::{self, ModelDataResolver};
 use crate::pb::type_proto::Value;
-use crate::pb::{self, TensorProto, TypeProto};
+use crate::pb::{self, GraphProto, OperatorSetIdProto, TensorProto, TypeProto};
 use crate::tensor::{load_tensor, translate_inference_fact};
 use prost::Message;
 
@@ -213,6 +213,60 @@ impl OnnxOpRegister {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct OnnxMetadata {
+    pub ir_version: i64,
+    pub opset_import: Vec<OperatorSetIdProto>,
+    pub producer_name: String,
+    pub producer_version: String,
+    pub domain: String,
+    pub model_version: i64,
+    pub doc_string: String,
+    pub graph: Option<GraphProto>,
+    pub metadata_props: HashMap<String,String> 
+}
+
+
+impl OnnxMetadata {
+    pub fn get_metadata(model_proto: &pb::ModelProto ) -> TractResult<OnnxMetadata> {
+        let parse_metadata_props: HashMap<String,String> = model_proto.to_owned().metadata_props
+            .into_iter().map(|entry| (entry.key, entry.value)).collect();
+       Ok(OnnxMetadata {
+            ir_version: model_proto.ir_version,
+            opset_import: model_proto.clone().opset_import,
+            producer_name: model_proto.clone().producer_name,
+            producer_version: model_proto.clone().producer_version,
+            domain: model_proto.clone().domain,
+            model_version: model_proto.model_version,
+            doc_string: model_proto.clone().doc_string,
+            graph: match model_proto.graph {
+                Some(_) => model_proto.clone().graph ,
+                None => None
+            },
+            metadata_props:  parse_metadata_props 
+        })
+    }
+}
+
+impl Default for OnnxMetadata {
+    fn default() -> Self {
+        let _opset = OperatorSetIdProto::default();
+        let graph = GraphProto::default();
+        OnnxMetadata {
+            ir_version: 0,
+            opset_import: vec![_opset],
+            producer_name: String::from(""),
+            producer_version: String::from(""),
+            domain: String::from(""),
+            model_version: 0,
+            doc_string: String::from(""),
+            graph: Some(graph),
+            metadata_props: HashMap::new()  
+        }
+    }
+}
+
+
 #[derive(Clone)]
 pub struct Onnx {
     pub op_register: OnnxOpRegister,
@@ -236,6 +290,23 @@ impl Onnx {
     pub fn parse(&self, proto: &pb::ModelProto, path: Option<&str>) -> TractResult<ParseResult> {
         self.parse_with_template(proto, path, Default::default())
     }
+
+    pub fn load_model_with_metadata(&mut self, model_path: impl AsRef<path::Path>) -> TractResult<(InferenceModel, OnnxMetadata)>{
+        let mut path = PathBuf::new();
+        path.push(&model_path);
+        let proto = self.proto_model_for_path(&model_path)?;
+        let mut dir: Option<&str> = None;
+        if let Some(dir_opt) = path.parent() {
+            dir = dir_opt.to_str();
+        }
+        let ParseResult { model, unresolved_inputs, .. } = self.parse(&proto, dir)?;
+        if unresolved_inputs.len() > 0 {
+            bail!("Could not resolve inputs at top-level: {:?}", unresolved_inputs)
+        }
+        let _metadata = OnnxMetadata::get_metadata(&proto)?; 
+        Ok((model, _metadata))
+    }
+
     pub fn parse_with_template(
         &self,
         proto: &pb::ModelProto,
@@ -248,8 +319,8 @@ impl Onnx {
             .find(|import| import.domain.is_empty() || import.domain == "ai.onnx")
             .map(|op| op.version)
             .unwrap_or(0);
-        let graph =
-            proto.graph.as_ref().ok_or_else(|| anyhow!("model proto does not contain a graph"))?;
+        // self.metadata = OnnxMetadata::get_metadata(&proto)?;
+        let graph = proto.graph.as_ref().ok_or_else(|| anyhow!("model proto does not contain a graph"))?;
         debug!("ONNX operator set version: {:?}", onnx_operator_set_version);
         if onnx_operator_set_version != 0 && !(9..19).contains(&onnx_operator_set_version) {
             warn!("ONNX operator for your model is {}, tract is only tested against \
@@ -267,7 +338,6 @@ impl Onnx {
         trace!("created ParsingContext");
         ctx.parse_graph(graph)
     }
-
     pub fn with_ignore_output_shapes(self, ignore: bool) -> Onnx {
         Self { use_output_shapes: !ignore, ..self }
     }