update: endpoint change

app.py

@@ -3,12 +3,11 @@ import netron
 import os
 import threading
 import time
+import requests
+import json
 from PIL import Image
 import cv2
 import numpy as np
-import torch
-from yolov5 import xai_yolov5
-from yolov8 import xai_yolov8s
 
 # Sample images directory
 sample_images = {
@@ -18,6 +17,7 @@ sample_images = {
 
 # Preloaded model file path (update this path as needed)
 preloaded_model_file = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")  # Example path
+onnx_endpoint = "http://localhost:8080/v1/models/yolov5:predict"  # ONNX model endpoint
 
 def load_sample_image(sample_name):
     """Load a sample image based on user selection."""
@@ -26,24 +26,40 @@ def load_sample_image(sample_name):
         return Image.open(image_path)
     return None
 
+def preprocess_image(image):
+    """Preprocess the image for ONNX model input."""
+    image = np.array(image)
+    image = cv2.resize(image, (640, 640))
+    _, buffer = cv2.imencode('.jpg', image)
+    image_bytes = buffer.tobytes()
+    return image_bytes
+
 def process_image(sample_choice, uploaded_image, yolo_versions):
-    """Process the image using selected YOLO models."""
+    """Process the image using selected YOLO models or the ONNX endpoint."""
     if uploaded_image is not None:
         image = uploaded_image  # Use the uploaded image
     else:
         image = load_sample_image(sample_choice)  # Use selected sample image
 
-    image = np.array(image)
-    image = cv2.resize(image, (640, 640))
+    image_bytes = preprocess_image(image)
     result_images = []
 
     for yolo_version in yolo_versions:
         if yolo_version == "yolov5":
-            result_images.append(xai_yolov5(image)) 
-        elif yolo_version == "yolov8s":
-            result_images.append(xai_yolov8s(image))
+            # Call the ONNX endpoint
+            headers = {"Content-Type": "application/json"}
+            data = {"inputs": [image_bytes.tolist()]}
+            try:
+                response = requests.post(onnx_endpoint, headers=headers, data=json.dumps(data))
+                if response.status_code == 200:
+                    results = response.json()
+                    result_images.append((image, str(results)))  # Example placeholder result
+                else:
+                    result_images.append((image, f"Error: {response.status_code}"))
+            except Exception as e:
+                result_images.append((image, f"Failed: {str(e)}"))
         else:
-            result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
+            result_images.append((image, f"{yolo_version} not yet implemented."))
 
     return result_images
 
@@ -52,6 +68,7 @@ def serve_netron(model_file):
     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
+
 def view_model():
     """Handle model visualization using preloaded model file."""
     if not os.path.exists(preloaded_model_file):
@@ -92,7 +109,7 @@ with gr.Blocks(css=custom_css) as interface:
             )
 
             selected_models = gr.CheckboxGroup(
-                choices=["yolov5", "yolov8s"],
+                choices=["yolov5"],
                 value=["yolov5"],
                 label="Select Model(s)",
             )