Spaces:
Sleeping
Sleeping
File size: 1,830 Bytes
45e6c67 07c1b75 3012224 45e6c67 ba0e7f4 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 |
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,):
pred = model(image, device=device, retina_masks=True, imgsz=1024, conf=conf_thres, iou=iou_thres)
prompt_process = FastSAMPrompt(input, pred, device="cpu")
ann = prompt_process.everything_prompt()
prompt_process.plot(annotations=ann, output_path="./output.jpg", withContours=False, better_quality=False)
output = Image.open('./output.jpg')
output = np.array(output)
return output
title = "FAST-SAM Segment Anything"
description = "A simple Gradio interface to infer on FAST-SAM model"
examples = [["image_1.jpg", 0.25, 0.45],
["image_2.jpg", 0.25, 0.45],
["image_3.jpg", 0.25, 0.45],
["image_4.jpg", 0.25, 0.45],
["image_5.jpg", 0.25, 0.45],
["image_6.jpg", 0.25, 0.45],
["image_7.jpg", 0.25, 0.45],
["image_8.jpg", 0.25, 0.45],
["image_9.jpg", 0.25, 0.45],
["image_10.jpg", 0.25, 0.45]]
demo = gr.Interface(inference,
inputs = [gr.Image(width=320, height=320, label="Input Image"),
gr.Slider(0, 1, 0.25, label="Confidence Threshold"),
gr.Slider(0, 1, 0.45, label="IoU Thresold")],
outputs= [gr.Image(width=640, height=640, label="Output")],
title=title,
description=description,
examples=examples)
demo.launch()
|