Update app.py

app.py

@@ -159,29 +159,22 @@ if __name__ == "__main__":
     interface.launch(share=True)
 """
 
-
 import gradio as gr
-import torch
-from ultralytics import YOLO
+import onnxruntime
 import os
 
-# Assuming the YOLOv5 model is located in the same directory as this script
-model_path = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")  # Replace with the actual model directory name
-
-def visualize_yolov5():
+def visualize_onnx_model(onnx_model_path):
   """
-  Visualizes the given YOLOv5 model using Netron.
+  Visualizes the given ONNX model using Netron.
+
+  Args:
+    onnx_model_path: The path to the ONNX model file.
   """
   try:
-    # Load the YOLOv5 model
-    model = YOLO(model_path) 
-
-    # Extract the PyTorch model 
-    pytorch_model = model.model 
-
-    # Save the PyTorch model to a temporary file
-    temp_model_path = "temp_model" 
-    torch.save(pytorch_model.state_dict(), os.path.join(temp_model_path, "pytorch_model.bin"))
+    # Save the ONNX model to a temporary file (optional, but can be helpful)
+    temp_model_path = "temp_model.onnx"
+    with open(temp_model_path, "wb") as f:
+      f.write(onnx_model_path.read())  # Assuming onnx_model_path is a file-like object
 
     # Run Netron 
     os.system(f"netron {temp_model_path}") 
@@ -191,12 +184,12 @@ def visualize_yolov5():
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=visualize_yolov5,
-    inputs=None,  # No input required
+    fn=visualize_onnx_model,
+    inputs="file",  # Accept the ONNX model as a file upload
     outputs="text", 
-    title="Netron YOLOv5 Model Visualization",
-    description="Visualize the YOLOv5 model."
+    title="Netron ONNX Model Visualization",
+    description="Upload an ONNX model file to visualize with Netron."
 )
 
-# Launch the Gradio app with public link
+# Launch the Gradio app
 iface.launch(share=True)