Spaces:
Sleeping
Sleeping
File size: 2,070 Bytes
d77b896 31e9d90 d77b896 e28fa28 d77b896 e28fa28 d77b896 e28fa28 d77b896 e28fa28 d77b896 e28fa28 d77b896 e28fa28 d77b896 31e9d90 d77b896 e28fa28 31e9d90 d77b896 e28fa28 d77b896 31e9d90 d77b896 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import gradio as gr
import onnx
from collections import Counter
import zipfile
import os
def process_onnx(uploaded_file):
# Check if the uploaded file is a zip 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] # Assuming there is only one ONNX file in the zip
file_path = os.path.join("/tmp", onnx_file)
else:
file_path = uploaded_file.name
# Load the ONNX model
model = onnx.load(file_path)
# Collect basic information about the model
info = {
"Model Name": model.graph.name,
"Number of Nodes": len(model.graph.node),
"Architecture Summary": Counter(),
"Nodes": []
}
# Iterate over each node (layer) to collect detailed information
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
# Format the summary output
summary_output = '\n'.join([f"{key}: {value}" for key, value in info.items() if key != "Nodes"])
# Format the complete nodes output
nodes_output = "Complete Nodes:\n" + '\n'.join([str(node) for node in info["Nodes"]])
return summary_output, nodes_output
# Define the Gradio Interface
iface = gr.Interface(
fn=process_onnx,
inputs=gr.File(label="Upload .ONNX or .ZIP File"),
outputs=[
gr.components.Textbox(label="Summary"),
gr.components.Textbox(label="Complete Nodes")
],
examples=["example1.onnx"], # Add your example file here
title="ONNX Model Scope",
description="Upload an ONNX file or a ZIP file 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."
)
# Launch the Interface
iface.launch(debug=True)
|