update: netron embedding

app.py

@@ -1,48 +1,38 @@
 import gradio as gr
-import netron
 import os
-import threading
 import time
-from PIL import Image
 import cv2
 import numpy as np
-import torch
-import base64
+from PIL import Image
 from yolov5 import xai_yolov5
 from yolov8 import xai_yolov8s
 
-# Sample images directory
+# Paths
 sample_images = {
     "Sample 1": os.path.join(os.getcwd(), "data/xai/sample1.jpeg"),
     "Sample 2": os.path.join(os.getcwd(), "data/xai/sample2.jpg"),
 }
+preloaded_model_file = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")  # Update as needed
+netron_html_file = os.path.join(os.getcwd(), "model_visualization.html")  # Netron exported file
 
-# Preloaded model file path (update this path as needed)
-preloaded_model_file = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")  # Example path
-
+# Load sample images
 def load_sample_image(sample_name):
-    """Load a sample image based on user selection."""
     image_path = sample_images.get(sample_name)
     if image_path and os.path.exists(image_path):
         return Image.open(image_path)
     return None
 
+# Process image with YOLO models
 def process_image(sample_choice, uploaded_image, yolo_versions):
-    """Process the image using selected YOLO models."""
     if uploaded_image is not None:
-        image = uploaded_image  # Use the uploaded image
+        image = uploaded_image
     else:
-        image = load_sample_image(sample_choice)  # Use selected sample image
+        image = load_sample_image(sample_choice)
 
     image = np.array(image)
     image = cv2.resize(image, (640, 640))
     result_images = []
 
-    # Encode image to base64
-    _, buffer = cv2.imencode('.jpg', image)
-    image_base64 = base64.b64encode(buffer).decode('utf-8')
-
-    # Process image with each selected YOLO version
     for yolo_version in yolo_versions:
         if yolo_version == "yolov5":
             result_images.append(xai_yolov5(image)) 
@@ -53,19 +43,14 @@ def process_image(sample_choice, uploaded_image, yolo_versions):
 
     return result_images
 
-def serve_netron(model_file):
-    """Start the Netron server in a separate thread."""
-    threading.Thread(target=netron.start, args=(model_file,), daemon=True).start()
-    time.sleep(1)  # Give some time for the server to start
-    return "http://localhost:8080"  # Default Netron URL
-
+# Embed Netron visualization (Static HTML Export)
 def view_model():
-    """Handle model visualization using preloaded model file."""
-    if not os.path.exists(preloaded_model_file):
-        return "Model file not found."
-    
-    netron_url = serve_netron(preloaded_model_file)
-    return f'<iframe src="{netron_url}" width="100%" height="600px"></iframe>'
+    if os.path.exists(netron_html_file):
+        with open(netron_html_file, "r") as f:
+            html_content = f.read()
+        return html_content
+    else:
+        return "<p style='color: red;'>Netron export not found. Generate HTML export and place it in the correct path.</p>"
 
 # Custom CSS for styling (optional)
 custom_css = """
@@ -78,13 +63,13 @@ custom_css = """
 }
 """
 
+# Gradio Interface
 with gr.Blocks(css=custom_css) as interface:
-    gr.Markdown("# XAI: Visualize Object Detection of Your Models")
-    
+    gr.Markdown("# XAI: Visualize Object Detection with YOLO Models")
+
     default_sample = "Sample 1"
 
     with gr.Row():
-        # Left side: Sample selection and upload image
         with gr.Column():
             sample_selection = gr.Radio(
                 choices=list(sample_images.keys()),
@@ -112,7 +97,7 @@ with gr.Blocks(css=custom_css) as interface:
                 label="Selected Sample Image",
             )
 
-    # Below the sample image, display results and architecture side by side
+    # Results and Netron Visualization
     with gr.Row():
         result_gallery = gr.Gallery(
             label="Results",
@@ -120,9 +105,9 @@ with gr.Blocks(css=custom_css) as interface:
             rows=1,
             height=500,
         )
+        netron_display = gr.HTML(label="Model Architecture")
 
-        netron_display = gr.HTML(label="Netron Visualization")
-
+    # Callbacks
     sample_selection.change(
         fn=load_sample_image,
         inputs=sample_selection,
@@ -135,9 +120,9 @@ with gr.Blocks(css=custom_css) as interface:
         outputs=[result_gallery],
     )
 
-    # Update Netron display when the interface loads
-    netron_display.value = view_model()  # Directly set the value
+    # Load Netron HTML into the iframe
+    netron_display.value = view_model()
 
-# Launching Gradio app and handling Netron visualization separately.
+# Launch Gradio
 if __name__ == "__main__":
     interface.launch(share=True)