Spaces:
Runtime error
Runtime error
File size: 1,848 Bytes
1d3f775 da2fab4 1d3f775 305e358 1d3f775 305e358 1d3f775 |
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 |
import gradio as gr
import numpy as np
import cv2
from PIL import Image
from ultralytics import YOLO
# Define available YOLO models
available_models = {
"X-ray": YOLO("xray.pt"),
"CT scan": YOLO("CT.pt"),
"Ultrasound": YOLO("ultrasound.pt"),
# Add more models as needed
}
# Create a function to perform image segmentation using the selected model
def segment_image(input_image, selected_model):
# Resize the input image to 255x255
img = np.array(input_image)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Perform object detection and segmentation using the selected model
model = available_models[selected_model]
results = model(img)
mask = results[0].masks.data.numpy()
target_height = img.shape[0]
target_width = img.shape[1]
# Resize the mask using OpenCV
resized_mask = cv2.resize(mask[0], (target_width, target_height))
resized_mask = (resized_mask * 255).astype(np.uint8)
# Create a copy of the original image
overlay_image = img.copy()
# Apply the resized mask to the overlay image
overlay_image[resized_mask > 0] = [100, 0, 0] # Overlay in green
# Convert the overlay image to PIL format
overlay_pil = Image.fromarray(overlay_image)
return overlay_pil
# Create the Gradio interface with a dropdown for model selection
iface = gr.Interface(
fn=segment_image,
inputs=[
gr.inputs.Image(type="pil", label="Upload an image"),
gr.inputs.Dropdown(
choices=list(available_models.keys()),
label="Select YOLO Model",
default="X-ray"
)
],
outputs=gr.outputs.Image(type="numpy", label="Segmented Image"),
title="YOLOv8 with SAM 😃",
description='This software generates the segmentation mask for Aorta for Point of Care Ultrasound (POCUS) images'
)
iface.launch()
|