import gradio as gr
import onnx
from collections import Counter
import zipfile
import os

def process_onnx(uploaded_file):
    
    if zipfile.is_zipfile(uploaded_file.name):
        with zipfile.ZipFile(uploaded_file.name, 'r') as zip_ref:
            zip_ref.extractall("/tmp")
            onnx_file = zip_ref.namelist()[0]  
            file_path = os.path.join("/tmp", onnx_file)
    else:
        file_path = uploaded_file.name

    
    model = onnx.load(file_path)

    
    info = {
        "Model Name": model.graph.name,
        "Number of Nodes": len(model.graph.node),
        "Architecture Summary": Counter(),
        "Nodes": []
    }

    
    for node in model.graph.node:
        node_info = {
            "Name": node.name,
            "Type": node.op_type,
            "Inputs": node.input,
            "Outputs": node.output
        }
        info["Nodes"].append(node_info)
        info["Architecture Summary"][node.op_type] += 1

    
    output = []
    for key, value in info.items():
        if key != "Nodes":
            output.append(f"{key}: {value}")
    output.append("\nComplete Nodes:")
    for node in info["Nodes"]:
        output.append(str(node))

    return '\n'.join(output)

iface = gr.Interface(
    fn=process_onnx,
    inputs=gr.File(label="Upload .ONNX or .ZIP (with .onnx) File"),
    outputs="text",
    title="ONNX Model Scope",
    description="Upload an ONNX file or a ZIP (containing an .onnx file) to extract and display its detailed information. This process can take some time depending on the size of the ONNX model."
)

iface.launch(debug=True)