Spaces:
Sleeping
Sleeping
File size: 2,998 Bytes
dfdcd97 a3ee867 e9cd6fd c95f3e0 9a34a8b e0d4d2f c95f3e0 9a34a8b 3cd1243 9a34a8b e0d4d2f 9a34a8b 2dd8fe8 26c0f04 3cd1243 73989e5 26c0f04 2dd8fe8 26c0f04 9a34a8b 3cd1243 9a34a8b 73989e5 9a34a8b 73989e5 2dd8fe8 9a34a8b 73989e5 3cd1243 3ba1061 73989e5 e0d4d2f e9cd6fd e0d4d2f e9cd6fd 3ba1061 3cd1243 e9cd6fd 3ba1061 9a34a8b 3ba1061 9a34a8b 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 78 79 80 81 82 83 84 85 |
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(annotations, prompt_process, mask_random_color=True, better_quality=True, retina=True, with_contours=True):
# ... (keep the existing plot function as is)
# This function doesn't need modification for our purposes
def segment_image(input_image, object_name):
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 with adjusted parameters
everything_results = model(input_image, retina_masks=True, imgsz=1024, conf=0.25, iou=0.7)
# Prepare a Prompt Process object
prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
# Use text prompt to segment the specified object
results = prompt_process.text_prompt(text=object_name)
if not results:
return input_image, f"Could not find '{object_name}' in the image."
# Post-process the masks
for ann in results:
if ann.masks is not None:
masks = ann.masks.data
if isinstance(masks[0], torch.Tensor):
masks = np.array(masks.cpu())
for i, mask in enumerate(masks):
# Apply more aggressive morphological operations
kernel = np.ones((5,5), np.uint8)
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
masks[i] = cv2.dilate(mask, kernel, iterations=2)
ann.masks.data = masks
# Plot the results
result_image = plot(annotations=results, prompt_process=prompt_process)
return result_image, f"Segmented '{object_name}' in the image."
except Exception as e:
return None, f"An error occurred: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=segment_image,
inputs=[
gr.Image(type="numpy", label="Upload an image"),
gr.Textbox(label="Specify object to segment (e.g., dog, cat, grass)")
],
outputs=[
gr.Image(type="pil", label="Segmented Image"),
gr.Textbox(label="Status")
],
title="FastSAM Segmentation with Object Specification",
description="Upload an image and specify an object to segment using FastSAM."
)
# Launch the interface
iface.launch() |