|
|
import torch |
|
|
import numpy as np |
|
|
import cv2 |
|
|
import gradio as gr |
|
|
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth").to("cuda") |
|
|
mask_generator = SamAutomaticMaskGenerator(sam) |
|
|
|
|
|
def segment(image): |
|
|
image = np.array(image) |
|
|
masks = mask_generator.generate(image) |
|
|
largest_mask = max(masks, key=lambda x: x['area'])['segmentation'] |
|
|
binary_mask = np.where(largest_mask, 255, 0).astype(np.uint8) |
|
|
return Image.fromarray(binary_mask) |
|
|
|
|
|
|
|
|
demo = gr.Interface(fn=segment, inputs="image", outputs="image") |
|
|
demo.launch() |
|
|
|
