Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| # Load YOLO model | |
| model = YOLO("best.pt") | |
| def segment_image(input_image): | |
| # Resize the input image to 255x255 | |
| img = np.array(input_image) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| # Perform object detection and segmentation | |
| # Perform object detection and segmentation | |
| 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() | |
| print(overlay_image.shape, mask.shape ) | |
| # 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 | |
| iface = gr.Interface( | |
| fn=segment_image, | |
| inputs=gr.components.Image(type="pil", label="Upload an image"), | |
| outputs=gr.components.Image(type="numpy", label="Segmented Image"), | |
| title="Aorta segmentation and Detection using YOLOv8 😃", | |
| description= 'This software generates the segementation mask for Aorta for the Point of Care Ultrasound (POCUS) images') | |
| iface.launch() |