Spaces:
Sleeping
Sleeping
File size: 2,304 Bytes
dfdcd97 a3ee867 e9cd6fd c95f3e0 9a34a8b e0d4d2f c95f3e0 9a34a8b 3cd1243 9a34a8b e0d4d2f 4f39124 72f4c5c 4f39124 72f4c5c 4f39124 72f4c5c 4f39124 73989e5 26c0f04 4f39124 26c0f04 9a34a8b 3cd1243 4f39124 2dd8fe8 9a34a8b 4f39124 73989e5 4f39124 3ba1061 73989e5 e0d4d2f e9cd6fd 4f39124 e9cd6fd 4f39124 e9cd6fd 3ba1061 9a34a8b 3ba1061 4f39124 e0d4d2f e9cd6fd |
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 68 69 70 71 72 73 74 75 76 77 |
import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import io
from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPrompt
# Set up device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load FastSAM model
model = FastSAM("FastSAM-s.pt") # or FastSAM-x.pt
def fig2img(fig):
buf = io.BytesIO()
fig.savefig(buf)
buf.seek(0)
img = Image.open(buf)
return img
def plot_masks(annotations, output_shape):
fig, ax = plt.subplots(figsize=(10, 10))
ax.imshow(annotations[0].orig_img)
for ann in annotations:
for mask in ann.masks.data:
mask = cv2.resize(mask.cpu().numpy().astype('uint8'), output_shape[::-1])
masked = np.ma.masked_where(mask == 0, mask)
ax.imshow(masked, alpha=0.5, cmap=plt.cm.get_cmap('jet'))
ax.axis('off')
plt.close()
return fig2img(fig)
def segment_everything(input_image):
try:
if input_image is None:
return None, "Please upload an image before submitting."
input_image = Image.fromarray(input_image).convert("RGB")
# Run FastSAM model in "everything" mode
everything_results = model(input_image, device=device, retina_masks=True, imgsz=1024, conf=0.25, iou=0.9, agnostic_nms=True)
# Prepare a Prompt Process object
prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
# Get everything segmentation
ann = prompt_process.everything_prompt()
# Plot the results
result_image = plot_masks(ann, input_image.size)
return result_image, f"Segmented everything in the image. Found {len(ann[0].masks)} objects."
except Exception as e:
return None, f"An error occurred: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=segment_everything,
inputs=[
gr.Image(type="numpy", label="Upload an image")
],
outputs=[
gr.Image(type="pil", label="Segmented Image"),
gr.Textbox(label="Status")
],
title="FastSAM Everything Segmentation",
description="Upload an image to segment all objects using FastSAM."
)
# Launch the interface
iface.launch() |