Spaces:
Runtime error
Runtime error
File size: 1,163 Bytes
aa9a932 |
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 |
import gradio as gr
import torch
import numpy as np
from PIL import Image
import cv2
from segment_anything import sam_model_registry, SamPredictor
# Load model
checkpoint = "sam_vit_h_4b8939.pth"
device = "cuda" if torch.cuda.is_available() else "cpu"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=checkpoint)
sam.to(device)
predictor = SamPredictor(sam)
def segment_image(input_img):
np_img = np.array(input_img)
image = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
predictor.set_image(image)
h, w, _ = image.shape
input_point = np.array([[w // 2, h // 2]])
input_label = np.array([1])
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=False
)
mask = masks[0].astype(np.uint8) * 255
return Image.fromarray(mask)
# UI
iface = gr.Interface(fn=segment_image,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
title="Segment Anything Model",
description="Upload an image and get a segmentation mask.")
iface.launch()
|