Update app.py

app.py

@@ -159,42 +159,39 @@ if __name__ == "__main__":
     interface.launch(share=True)
 """
 
+
 import gradio as gr
 import netron
 import os
 import threading
-from flask import Flask, send_from_directory
-
-# Flask app to serve the Netron model visualization
-app = Flask(__name__)
 
-def start_netron(model_path):
-    # Start a new thread to serve the Netron visualization on a local server
-    netron.start(model_path, port=8080)
-
-# Define a function to visualize the model and generate an iframe
+# Define the function to visualize the model
 def visualize_model(model_file):
-    # Save the model file temporarily
-    temp_model_path = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")
-    with open(temp_model_path, "wb") as f:
-        f.write(model_file.read())
+    # If the model file is a path (file uploaded), read it
+    if isinstance(model_file, str):
+        model_path = model_file
+    else:
+        # If it's a NamedString (new behavior), save it to a temporary file
+        temp_model_path = f"./temp_model.{model_file.name.split('.')[-1]}"
+        with open(temp_model_path, "wb") as f:
+            f.write(model_file.encode())  # Use encode() to write the string content
+        model_path = temp_model_path
 
-    # Start Netron in a separate thread
-    threading.Thread(target=start_netron, args=(temp_model_path,)).start()
+    # Start Netron visualization in a separate thread
+    threading.Thread(target=netron.start, args=(model_path,)).start()
 
     # Return an iframe embedding the Netron visualization
     iframe_html = f'<iframe src="http://localhost:8080" width="800" height="600"></iframe>'
     return iframe_html
 
-# Gradio interface to upload the model and show the visualization
+# Create the Gradio interface
 interface = gr.Interface(
     fn=visualize_model,
-    inputs=gr.File(label="Upload Model File"),
-    outputs=gr.HTML(),  # Output is HTML to embed iframe
+    inputs=gr.File(label="Upload Model File"),  # Model file input
+    outputs=gr.HTML(),  # Output is HTML to embed the iframe
     live=True
 )
 
-# Start the Gradio interface
 interface.launch()