Fix: interface update

app.py

@@ -6,7 +6,7 @@ import torchvision.transforms as transforms
 import gradio as gr
 from yolov5 import xai_yolov5
 from yolov8 import xai_yolov8s
-
+"""
 def process_image(image, yolo_versions=["yolov5"]):
     image = np.array(image)
     image = cv2.resize(image, (640, 640))
@@ -20,7 +20,7 @@ def process_image(image, yolo_versions=["yolov5"]):
             result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
     return result_images
 
-
+"""
 sample_images = {
     "Sample 1": os.path.join(os.getcwd(), "data/xai/sample1.jpeg"),
     "Sample 2":  os.path.join(os.getcwd(), "data/xai/sample2.jpg"),
@@ -59,23 +59,41 @@ with gr.Blocks() as interface:
         outputs=result_gallery,
     )
 """
+def load_sample_image(choice):
+    if choice in sample_images:
+        image_path = sample_images[choice]
+        return cv2.imread(image_path)[:, :, ::-1]  # Convert BGR to RGB for display
+    else:
+        raise ValueError("Invalid sample selection.")
+
+
+def process_image(choice, yolo_versions=["yolov5"]):
+    image = load_sample_image(choice)
+    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)) 
+        elif yolo_version == "yolov8s":
+            result_images.append(xai_yolov8s(image))
+        else:
+            result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
+    return result_images
 
+
+import gradio as gr
 with gr.Blocks() as interface:
     gr.Markdown("# XAI: Visualize Object Detection of Your Models")
     gr.Markdown("Select a sample image to visualize object detection.")
-
-    with gr.Row():
-        sample_selection = gr.Radio(
-            choices=list(sample_images.keys()),
-            label="Select a Sample Image",
-            type="value",
-        )
-        sample_display = gr.Image(label="Sample Image Preview", value=default_sample_image)
-        sample_selection.change(fn=load_sample_image, inputs=sample_selection, outputs=sample_display)
-
+    sample_selection = gr.Radio(
+        choices=list(sample_images.keys()),
+        label="Select a Sample Image",
+        type="value",
+    )
     selected_models = gr.CheckboxGroup(
         choices=["yolov5", "yolov8s"],
-        value=["yolov5"], 
+        value=["yolov5"],
         label="Select Model(s)",
     )
     result_gallery = gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500)