Spaces:
Sleeping
Sleeping
File size: 2,543 Bytes
45e6c67 07c1b75 3012224 45e6c67 4dcd8cb 45e6c67 e93d8c1 45e6c67 4dcd8cb 45e6c67 4dcd8cb 45e6c67 4dcd8cb 45e6c67 4dcd8cb 45e6c67 |
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 |
import torch
import numpy as np
import gradio as gr
import torch.nn.functional as F
from PIL import Image
from fastsam import FastSAM, FastSAMPrompt
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = "mps"
model = FastSAM('./weights/FastSAM-x.pt')
model.to(device)
def inference(image, conf_thres, iou_thres, text):
pred = model(image, device=device, retina_masks=True, imgsz=1024, conf=conf_thres, iou=iou_thres)
prompt_process = FastSAMPrompt(image, pred, device="cpu")
ann = prompt_process.everything_prompt()
prompt_process.plot(annotations=ann, output_path="./output_sam.jpg", withContours=False, better_quality=False)
output_sam = Image.open('./output_sam.jpg')
output_sam = np.array(output_sam)
output_text = None
if text:
ann = prompt_process.text_prompt(text=text)
prompt_process.plot(annotations=ann, output_path="./output_text.jpg", withContours=False, better_quality=False)
output_text = Image.open('./output_text.jpg')
output_text = np.array(output_text)
return output_sam, output_text
title = "FAST-SAM Segment Anything"
description = "A simple Gradio interface to infer on FAST-SAM model"
examples = [["image_1.jpg", 0.25, 0.45, 'A black tire'],
["image_2.jpg", 0.25, 0.45, 'Shades of blue'],
["image_3.jpg", 0.25, 0.45, 'A spiral staircase'],
["image_4.jpg", 0.25, 0.45, 'A clock and a plane'],
["image_5.jpg", 0.25, 0.45, 'Clouds in the sky'],
["image_6.jpg", 0.25, 0.45, 'Front wheel'],
["image_7.jpg", 0.25, 0.45, 'A white chair'],
["image_8.jpg", 0.25, 0.45, 'The grassy field'],
["image_9.jpg", 0.25, 0.45, 'Rock formation'],
["image_10.jpg", 0.25, 0.45, 'A rope railing']]
demo = gr.Interface(inference,
inputs = [gr.Image(width=640, height=640, label="Input Image"),
gr.Slider(0, 1, 0.25, label="Confidence Threshold"),
gr.Slider(0, 1, 0.45, label="IoU Thresold"),
gr.Textbox(label="Enter text promp", type="text"),],
outputs= [gr.Image(width=640, height=640, label="Output SAM"),
gr.Image(width=640, height=640, label="Output text prompt")],
title=title,
description=description,
examples=examples)
demo.launch()
|