update: file encoding

app.py

@@ -3,11 +3,13 @@ import netron
 import os
 import threading
 import time
-import requests
-import json
 from PIL import Image
 import cv2
 import numpy as np
+import torch
+import base64
+from yolov5 import xai_yolov5
+from yolov8 import xai_yolov8s
 
 # Sample images directory
 sample_images = {
@@ -17,7 +19,6 @@ 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,40 +27,29 @@ 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 or the ONNX endpoint."""
+    """Process the image using selected YOLO models."""
     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_bytes = preprocess_image(image)
+    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":
-            # 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)}"))
+            result_images.append(xai_yolov5(image)) 
+        elif yolo_version == "yolov8s":
+            result_images.append(xai_yolov8s(image))
         else:
-            result_images.append((image, f"{yolo_version} not yet implemented."))
+            result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
 
     return result_images
 
@@ -109,7 +99,7 @@ with gr.Blocks(css=custom_css) as interface:
             )
 
             selected_models = gr.CheckboxGroup(
-                choices=["yolov5"],
+                choices=["yolov5", "yolov8s"],
                 value=["yolov5"],
                 label="Select Model(s)",
             )