update: dynamic html support

app.py

@@ -1,56 +1,68 @@
 import gradio as gr
+import netron
 import os
+import threading
 import time
+from PIL import Image
 import cv2
 import numpy as np
-from PIL import Image
 from yolov5 import xai_yolov5
 from yolov8 import xai_yolov8s
 
-# Paths
+# Sample images directory
 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
 
-# Load sample images
+# Preloaded model file path
+preloaded_model_file = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")  # Example path
+
 def load_sample_image(sample_name):
     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):
+    # Use uploaded or sample image
     if uploaded_image is not None:
         image = uploaded_image
     else:
         image = load_sample_image(sample_choice)
 
+    # Resize and process the image
     image = np.array(image)
     image = cv2.resize(image, (640, 640))
     result_images = []
 
     for yolo_version in yolo_versions:
         if yolo_version == "yolov5":
-            result_images.append(xai_yolov5(image)) 
+            result_images.append(xai_yolov5(image))
         elif yolo_version == "yolov8s":
             result_images.append(xai_yolov8s(image))
         else:
-            result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
+            result_images.append((Image.fromarray(image), f"{yolo_version} not implemented."))
 
     return result_images
 
-# Embed Netron visualization (Static HTML Export)
-def view_model():
-    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>"
+def start_netron_backend(model_file):
+    def serve_netron():
+        netron.start(model_file, port=8080)  # Start Netron on port 8080
+    
+    # Launch Netron in a separate thread
+    threading.Thread(target=serve_netron, daemon=True).start()
+    time.sleep(2)  # Allow server to initialize
+
+def view_netron_model():
+    # Ensure model exists
+    if not os.path.exists(preloaded_model_file):
+        return "Model file not found."
+
+    # Start Netron backend
+    start_netron_backend(preloaded_model_file)
+    return '<iframe src="http://localhost:8080/" width="100%" height="600px"></iframe>'
+
 
 # Custom CSS for styling (optional)
 custom_css = """
@@ -63,13 +75,15 @@ custom_css = """
 }
 """
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks(css=custom_css) as interface:
-    gr.Markdown("# XAI: Visualize Object Detection with YOLO Models")
+    gr.Markdown("# XAI: Visualize Object Detection of Your Models")
 
+    # Default sample
     default_sample = "Sample 1"
 
     with gr.Row():
+        # Left: Select sample or upload image
         with gr.Column():
             sample_selection = gr.Radio(
                 choices=list(sample_images.keys()),
@@ -78,10 +92,7 @@ with gr.Blocks(css=custom_css) as interface:
                 value=default_sample,
             )
 
-            upload_image = gr.Image(
-                label="Upload an Image",
-                type="pil",  
-            )
+            upload_image = gr.Image(label="Upload an Image", type="pil")
 
             selected_models = gr.CheckboxGroup(
                 choices=["yolov5", "yolov8s"],
@@ -91,13 +102,14 @@ with gr.Blocks(css=custom_css) as interface:
 
             run_button = gr.Button("Run", elem_id="run_button")
 
+        # Right: Display sample image
         with gr.Column():
             sample_display = gr.Image(
-                value=load_sample_image(default_sample),  
+                value=load_sample_image(default_sample),
                 label="Selected Sample Image",
             )
 
-    # Results and Netron Visualization
+    # Results and Netron
     with gr.Row():
         result_gallery = gr.Gallery(
             label="Results",
@@ -105,24 +117,26 @@ with gr.Blocks(css=custom_css) as interface:
             rows=1,
             height=500,
         )
-        netron_display = gr.HTML(label="Model Architecture")
 
-    # Callbacks
+        netron_display = gr.HTML(label="Model Visualization")
+
+    # Sample selection update
     sample_selection.change(
         fn=load_sample_image,
         inputs=sample_selection,
         outputs=sample_display,
     )
 
+    # Process image
     run_button.click(
         fn=process_image,
         inputs=[sample_selection, upload_image, selected_models],
         outputs=[result_gallery],
     )
 
-    # Load Netron HTML into the iframe
-    netron_display.value = view_model()
+    # Netron model visualization
+    netron_display.value = view_netron_model()
 
-# Launch Gradio
+# Launch Gradio app
 if __name__ == "__main__":
     interface.launch(share=True)