Spaces:
Sleeping
Sleeping
File size: 2,157 Bytes
60af537 57419d8 9b2c5e1 57419d8 d4cb7c6 57419d8 6c34a8c c5224aa 57419d8 b21eb38 6263059 e8f8ad5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import numpy as np
import cv2
import os
from PIL import Image
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))
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
# Paths to sample images
sample_images = [
os.path.join(os.getcwd(), "data/xai/sample1.jpeg"),
os.path.join(os.getcwd(), "data/xai/sample2.jpg")
]
def load_sample(sample_name):
if sample_name and sample_name in sample_images:
return Image.open(sample_images[sample_name])
return None
interface = gr.Interface(
fn=process_image,
inputs=[
gr.Image(type="pil", label="Upload an Image"),
gr.CheckboxGroup(
choices=["yolov3", "yolov8s"],
value=["yolov5"], # Set the default value (YOLOv5 checked by default)
label="Select Model(s)",
),
gr.Dropdown(
choices=list(sample_images.keys()),
label="Select a Sample Image",
type="value",
interactive=True,
),
],
outputs=gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500),
title="Visualising the key image features that drive decisions with our explainable AI tool.",
description="XAI: Upload an image or select a sample to visualize object detection of your models.",
)
def main_logic(uploaded_image, selected_models, sample_selection):
# If the user selects a sample image, use that instead of the uploaded one
if sample_selection:
image = load_sample(sample_selection)
else:
image = uploaded_image
# Call the processing function
return process_image(image, selected_models)
interface.launch()
|