rknncli 0.2.0__py3-none-any.whl

rknncli/__init__.py

@@ -0,0 +1,3 @@
+"""RKNN CLI - A command line tool for parsing and displaying RKNN model information."""
+
+__version__ = "0.1.0"

rknncli/cli.py

@@ -0,0 +1,229 @@
+"""RKNN CLI - Command line interface for parsing RKNN models."""
+
+import argparse
+import sys
+from pathlib import Path
+from typing import List, Dict, Any
+
+from rknncli.parser import RKNNParser
+
+
+def format_shape(size: List) -> List:
+    """Format tensor shape with dimension names.
+
+    Args:
+        size: List of dimension sizes.
+
+    Returns:
+        List of dimension names (strings) or dimension values.
+    """
+    return size
+
+
+def get_dtype_str(dtype_info: Dict) -> str:
+    """Get data type string from dtype info.
+
+    Args:
+        dtype_info: Dictionary containing dtype information.
+
+    Returns:
+        String representation of the data type.
+    """
+    if isinstance(dtype_info, dict):
+        # Try vx_type first, then qnt_type
+        vx_type = dtype_info.get("vx_type", "").strip()
+        if vx_type:
+            return vx_type.upper()
+        qnt_type = dtype_info.get("qnt_type", "").strip()
+        if qnt_type:
+            return qnt_type.upper()
+    return "FLOAT"
+
+
+
+
+def print_merged_model_info(parser: RKNNParser) -> None:
+    """Print merged model information from both FlatBuffers and JSON.
+
+    Args:
+        parser: RKNNParser instance with FlatBuffers support.
+    """
+    # Get basic model info from JSON
+    print(f"Model: {parser.get_model_name()}")
+    platforms = parser.get_target_platform()
+    if platforms:
+        print(f"Target Platform: {', '.join(platforms)}")
+
+    # Get FlatBuffers info
+    fb_info = parser.get_flatbuffers_info()
+    if fb_info:
+        print(f"Format: {fb_info.get('format', 'Unknown')}")
+        print(f"Source: {fb_info.get('source', 'Unknown')}")
+        print(f"Compiler: {fb_info.get('compiler', 'Unknown')}")
+        print(f"Runtime: {fb_info.get('runtime', 'Unknown')}")
+
+        if fb_info.get("num_graphs", 0) > 0:
+            print(f"Number of graphs: {fb_info['num_graphs']}")
+
+    print()
+
+
+def print_merged_io_info(parser: RKNNParser) -> None:
+    """Print merged input/output information from both FlatBuffers and JSON.
+
+    Args:
+        parser: RKNNParser instance with FlatBuffers support.
+    """
+    # Get merged IO info
+    inputs, outputs = parser.get_merged_io_info()
+
+    # Print input information
+    print("Input information")
+    print("-" * 80)
+
+    for tensor in inputs:
+        name = tensor.get("url", f"tensor_{tensor.get('tensor_id', 0)}")
+        dtype = tensor.get("dtype", {})
+
+        # Get dtype string
+        if isinstance(dtype, dict):
+            dtype_str = get_dtype_str(dtype)
+        else:
+            dtype_str = str(dtype).upper()
+
+        size = tensor.get("size", [])
+        shape = format_shape(size)
+        shape_str = "[" + ", ".join(f"'{s}'" if isinstance(s, str) else str(s) for s in shape) + "]"
+
+        # Build the base output string
+        output_parts = [f'  ValueInfo "{name}": type {dtype_str}, shape {shape_str}']
+
+        # Add layout info from FlatBuffers
+        if "layout" in tensor and tensor["layout"]:
+            layout = tensor["layout"].upper()
+            layout_ori = tensor.get("layout_ori", "").upper()
+            if layout_ori and layout_ori != layout:
+                output_parts.append(f"layout {layout}(ori:{layout_ori})")
+            else:
+                output_parts.append(f"layout {layout}")
+
+        # Add quantization info from FlatBuffers
+        if "quant_info" in tensor:
+            quant = tensor["quant_info"]
+            if quant.get("qmethod") or quant.get("qtype"):
+                output_parts.append(f"quant {quant['qmethod']} {quant['qtype']}")
+
+        # Print the merged info on one line
+        print(", ".join(output_parts) + ",")
+
+    print()
+
+    # Print output information
+    print("Output information")
+    print("-" * 80)
+
+    for tensor in outputs:
+        name = tensor.get("url", f"tensor_{tensor.get('tensor_id', 0)}")
+        dtype = tensor.get("dtype", {})
+
+        # Get dtype string
+        if isinstance(dtype, dict):
+            dtype_str = get_dtype_str(dtype)
+        else:
+            dtype_str = str(dtype).upper()
+
+        size = tensor.get("size", [])
+        shape = format_shape(size)
+        shape_str = "[" + ", ".join(f"'{s}'" if isinstance(s, str) else str(s) for s in shape) + "]"
+
+        # Build the base output string
+        output_parts = [f'  ValueInfo "{name}": type {dtype_str}, shape {shape_str}']
+
+        # Add layout info from FlatBuffers
+        if "layout" in tensor and tensor["layout"]:
+            layout = tensor["layout"].upper()
+            layout_ori = tensor.get("layout_ori", "").upper()
+            if layout_ori and layout_ori != layout:
+                output_parts.append(f"layout {layout}(ori:{layout_ori})")
+            else:
+                output_parts.append(f"layout {layout}")
+
+        # Add quantization info from FlatBuffers
+        if "quant_info" in tensor:
+            quant = tensor["quant_info"]
+            if quant.get("qmethod") or quant.get("qtype"):
+                output_parts.append(f"quant {quant['qmethod']} {quant['qtype']}")
+
+        # Print the merged info on one line
+        print(", ".join(output_parts) + ",")
+
+
+def print_model_summary(parser) -> None:
+    """Print model summary information.
+
+    Args:
+        parser: RKNNParser instance.
+    """
+    print(f"Model: {parser.get_model_name()}")
+    print(f"Version: {parser.get_version()}")
+    platforms = parser.get_target_platform()
+    if platforms:
+        print(f"Target Platform: {', '.join(platforms)}")
+    print()
+
+
+
+
+def main() -> int:
+    """Main entry point.
+
+    Returns:
+        Exit code (0 for success, non-zero for error).
+    """
+    parser = argparse.ArgumentParser(
+        prog="rknncli",
+        description="A command line tool for parsing and displaying RKNN model information.",
+    )
+    parser.add_argument(
+        "model",
+        type=str,
+        help="Path to the RKNN model file",
+    )
+    parser.add_argument(
+        "-v", "--version",
+        action="version",
+        version="%(prog)s 0.1.0",
+    )
+
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    if not model_path.exists():
+        print(f"Error: File not found: {model_path}", file=sys.stderr)
+        return 1
+
+    if not model_path.is_file():
+        print(f"Error: Not a file: {model_path}", file=sys.stderr)
+        return 1
+
+    try:
+        # Always parse both FlatBuffers and JSON data
+        rknn_parser = RKNNParser(model_path, parse_flatbuffers=True)
+
+        # Print merged model information
+        print_merged_model_info(rknn_parser)
+
+        # Print merged IO information
+        print_merged_io_info(rknn_parser)
+    except ValueError as e:
+        print(f"Error: Failed to parse RKNN file: {e}", file=sys.stderr)
+        return 1
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

rknncli/parser.py

@@ -0,0 +1,320 @@
+"""RKNN file parser with FlatBuffers support."""
+
+import json
+import struct
+from pathlib import Path
+from typing import Any, Optional, Dict, List, Union, Tuple
+
+import flatbuffers
+from rknncli.schema.rknn.Model import Model
+from rknncli.schema.rknn.Graph import Graph
+from rknncli.schema.rknn.Tensor import Tensor
+from rknncli.schema.rknn.Node import Node
+
+
+class RKNNParser:
+    """Parser for RKNN model files with optional FlatBuffers support."""
+
+    HEADER_SIZE = 64
+    MAGIC_NUMBER = b"RKNN"
+
+    def __init__(self, file_path: Union[str, Path], parse_flatbuffers: bool = True):
+        """Initialize parser with RKNN file path.
+
+        Args:
+            file_path: Path to the RKNN model file.
+            parse_flatbuffers: Whether to parse FlatBuffers data. Defaults to False.
+        """
+        self.file_path = Path(file_path)
+        self.header: Dict[str, Any] = {}
+        self.model_info: Dict[str, Any] = {}
+        self.fb_model: Optional[Model] = None
+        self._parse(parse_flatbuffers)
+
+    def _parse(self, parse_flatbuffers: bool) -> None:
+        """Parse the RKNN file.
+
+        Args:
+            parse_flatbuffers: Whether to parse FlatBuffers data.
+        """
+        with open(self.file_path, "rb") as f:
+            # Read header
+            header_data = f.read(self.HEADER_SIZE)
+            if len(header_data) < self.HEADER_SIZE:
+                raise ValueError(f"File too small: {self.file_path}")
+
+            # Parse header
+            # Bytes 0-3: Magic number "RKNN"
+            magic = header_data[0:4]
+            if magic != self.MAGIC_NUMBER:
+                raise ValueError(f"Invalid magic number: {magic!r}, expected {self.MAGIC_NUMBER!r}")
+
+            # Bytes 4-7: Padding (4 bytes, zeros)
+            padding = header_data[4:8]
+
+            # Bytes 8-15: File format version (8 bytes, little-endian uint64)
+            file_format = struct.unpack("<Q", header_data[8:16])[0]
+
+            # Bytes 16-23: File length (8 bytes, little-endian uint64)
+            file_length = struct.unpack("<Q", header_data[16:24])[0]
+
+            self.header = {
+                "magic": magic,
+                "padding": padding,
+                "file_format": file_format,
+                "file_length": file_length,
+            }
+
+            # Calculate JSON offset and size
+            real_header_size = self.HEADER_SIZE
+            if file_format <= 1:
+                # only 3 uint64 for rknn-v1
+                real_header_size = 24
+
+            # Parse FlatBuffers data if requested
+            if parse_flatbuffers and file_length > 0:
+                fb_offset = real_header_size
+                fb_data = f.read(file_length)
+                if len(fb_data) == file_length:
+                    self.fb_model = Model.GetRootAs(fb_data, 0)
+
+            # Read JSON model info
+            file_size = self.file_path.stat().st_size
+            json_offset = real_header_size + file_length
+            f.seek(json_offset)
+            json_size = struct.unpack("<Q", f.read(8))[0]
+            if (
+                json_size <= 0 or
+                json_size > file_size - real_header_size - file_length - 8
+            ):
+                raise ValueError(f"Invalid JSON size: {json_size}")
+
+            json_data = f.read(json_size)
+
+            try:
+                self.model_info = json.loads(json_data.decode("utf-8"))
+            except (json.JSONDecodeError, UnicodeDecodeError) as e:
+                raise ValueError(f"Failed to parse model info JSON: {e}")
+
+    def get_flatbuffers_info(self) -> Dict[str, Any]:
+        """Get FlatBuffers model information.
+
+        Returns:
+            Dictionary containing FlatBuffers model data.
+        """
+        if not self.fb_model:
+            return {}
+
+        info = {}
+
+        # Basic model info
+        if self.fb_model.Format():
+            info["format"] = self.fb_model.Format().decode('utf-8')
+
+        if self.fb_model.Generator():
+            info["generator"] = self.fb_model.Generator().decode('utf-8')
+
+        if self.fb_model.Compiler():
+            info["compiler"] = self.fb_model.Compiler().decode('utf-8')
+
+        if self.fb_model.Runtime():
+            info["runtime"] = self.fb_model.Runtime().decode('utf-8')
+
+        if self.fb_model.Source():
+            info["source"] = self.fb_model.Source().decode('utf-8')
+
+        # Graphs info
+        info["num_graphs"] = self.fb_model.GraphsLength()
+
+        # Input/Output JSON strings
+        if self.fb_model.InputJson():
+            info["input_json"] = self.fb_model.InputJson().decode('utf-8')
+
+        if self.fb_model.OutputJson():
+            info["output_json"] = self.fb_model.OutputJson().decode('utf-8')
+
+        # Parse graphs
+        graphs = []
+        for i in range(self.fb_model.GraphsLength()):
+            graph = self.fb_model.Graphs(i)
+            if graph:
+                graph_info = {
+                    "num_tensors": graph.TensorsLength(),
+                    "num_nodes": graph.NodesLength(),
+                    "num_inputs": graph.InputsLength(),
+                    "num_outputs": graph.OutputsLength(),
+                }
+
+                # Parse tensors
+                tensors = []
+                for j in range(graph.TensorsLength()):
+                    tensor = graph.Tensors(j)
+                    if tensor:
+                        tensor_info = {
+                            "data_type": tensor.DataType(),
+                            "kind": tensor.Kind(),
+                            "name": tensor.Name().decode('utf-8') if tensor.Name() else "",
+                            "shape": [tensor.Shape(k) for k in range(tensor.ShapeLength())],
+                            "size": tensor.Size(),
+                            "index": tensor.Index(),
+                        }
+                        tensors.append(tensor_info)
+                graph_info["tensors"] = tensors
+
+                # Parse nodes
+                nodes = []
+                for j in range(graph.NodesLength()):
+                    node = graph.Nodes(j)
+                    if node:
+                        node_info = {
+                            "type": node.Type().decode('utf-8') if node.Type() else "",
+                            "name": node.Name().decode('utf-8') if node.Name() else "",
+                            "num_inputs": node.InputsLength(),
+                            "num_outputs": node.OutputsLength(),
+                        }
+                        nodes.append(node_info)
+                graph_info["nodes"] = nodes
+
+                graphs.append(graph_info)
+
+        info["graphs"] = graphs
+
+        return info
+
+    def get_input_info(self) -> List[Dict[str, Any]]:
+        """Get model input information.
+
+        Returns:
+            List of input tensor information.
+        """
+        inputs = []
+        norm_tensors = {t["tensor_id"]: t for t in self.model_info.get("norm_tensor", [])}
+
+        # Find input tensors from graph connections
+        for conn in self.model_info.get("graph", []):
+            if conn.get("left") == "input":
+                tensor_id = conn.get("right_tensor_id")
+                if tensor_id in norm_tensors:
+                    inputs.append(norm_tensors[tensor_id])
+
+        return inputs
+
+    def get_output_info(self) -> List[Dict[str, Any]]:
+        """Get model output information.
+
+        Returns:
+            List of output tensor information.
+        """
+        outputs = []
+        norm_tensors = {t["tensor_id"]: t for t in self.model_info.get("norm_tensor", [])}
+
+        # Find output tensors from graph connections
+        for conn in self.model_info.get("graph", []):
+            if conn.get("left") == "output":
+                tensor_id = conn.get("right_tensor_id")
+                if tensor_id in norm_tensors:
+                    outputs.append(norm_tensors[tensor_id])
+
+        return outputs
+
+    def get_model_name(self) -> str:
+        """Get model name."""
+        return self.model_info.get("name", "Unknown")
+
+    def get_version(self) -> str:
+        """Get model version."""
+        return self.model_info.get("version", "Unknown")
+
+    def get_target_platform(self) -> list[str]:
+        """Get target platform."""
+        return self.model_info.get("target_platform", [])
+
+    def get_merged_io_info(self) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+        """Get merged input/output information from both FlatBuffers and JSON.
+
+        Returns:
+            Tuple of (inputs, outputs) with merged information.
+        """
+        # Get basic IO info from JSON
+        json_inputs = self.get_input_info()
+        json_outputs = self.get_output_info()
+
+        # If no FlatBuffers data, return JSON data as is
+        if not self.fb_model:
+            return json_inputs, json_outputs
+
+        # Get FlatBuffers info
+        fb_info = self.get_flatbuffers_info()
+
+        # Extract generator attributes and quant info
+        generator_attrs = {}
+        quant_tab = {}
+
+        # Parse generator JSON string if present
+        if fb_info.get("generator"):
+            try:
+                # Use ast.literal_eval to handle single quotes
+                import ast
+                gen_data = ast.literal_eval(fb_info["generator"])
+                if isinstance(gen_data, dict):
+                    generator_attrs = gen_data.get("attrs", {})
+                    quant_tab = gen_data.get("quant_tab", {})
+            except (ValueError, SyntaxError):
+                # Fallback: try JSON parse after fixing quotes
+                try:
+                    import json
+                    fixed_str = fb_info["generator"].replace("'", '"')
+                    gen_data = json.loads(fixed_str)
+                    if isinstance(gen_data, dict):
+                        generator_attrs = gen_data.get("attrs", {})
+                        quant_tab = gen_data.get("quant_tab", {})
+                except (json.JSONDecodeError, ValueError):
+                    pass
+
+        # Merge input information
+        merged_inputs = []
+        for inp in json_inputs:
+            io_name = inp.get("url", "")
+            merged_inp = inp.copy()
+
+            # Add layout info from generator attrs
+            if io_name in generator_attrs:
+                merged_inp["layout"] = generator_attrs[io_name].get("layout", "")
+                merged_inp["layout_ori"] = generator_attrs[io_name].get("layout_ori", "")
+
+            # Add quant info from quant_tab
+            if io_name in quant_tab:
+                merged_inp["dtype"] = quant_tab[io_name].get("dtype", inp.get("dtype", {}))
+                merged_inp["quant_info"] = {
+                    "qmethod": quant_tab[io_name].get("qmethod", ""),
+                    "qtype": quant_tab[io_name].get("qtype", ""),
+                    "scale": quant_tab[io_name].get("scale", []),
+                    "zero_point": quant_tab[io_name].get("zero_point", [])
+                }
+
+            merged_inputs.append(merged_inp)
+
+        # Merge output information
+        merged_outputs = []
+        for out in json_outputs:
+            io_name = out.get("url", "")
+            merged_out = out.copy()
+
+            # Add layout info from generator attrs
+            if io_name in generator_attrs:
+                merged_out["layout"] = generator_attrs[io_name].get("layout", "")
+                merged_out["layout_ori"] = generator_attrs[io_name].get("layout_ori", "")
+
+            # Add quant info from quant_tab
+            if io_name in quant_tab:
+                merged_out["dtype"] = quant_tab[io_name].get("dtype", out.get("dtype", {}))
+                merged_out["quant_info"] = {
+                    "qmethod": quant_tab[io_name].get("qmethod", ""),
+                    "qtype": quant_tab[io_name].get("qtype", ""),
+                    "scale": quant_tab[io_name].get("scale", []),
+                    "zero_point": quant_tab[io_name].get("zero_point", [])
+                }
+
+            merged_outputs.append(merged_out)
+
+        return merged_inputs, merged_outputs