sam_everything.py

from segment_anything import SamPredictor, sam_model_registry,SamAutomaticMaskGenerator
from PIL import Image
import torch
from detectron2.data.detection_utils import read_image,pil_image_to_numpy
from detectron2.utils.visualizer import Visualizer
import numpy as np
from skimage import measure
import threading

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def seg_with_promp(imput_image,point_coords=None,box=None):
    if isinstance(imput_image, Image.Image):
        imput_image = pil_image_to_numpy(imput_image)
    point_labels = None
    if point_coords is not None:
        point_labels = np.ones(point_coords.shape[0])
    sam = sam_model_registry["vit_b"](checkpoint="./sam_vit_b_01ec64.pth").to(device)
    predictor = SamPredictor(sam)
    predictor.set_image(imput_image)

    masks = None

    if box is not None:
        masks, _, _ = predictor.predict(box=box)
    elif point_coords is not None and point_labels is not None:
        masks, _, _ = predictor.predict(point_coords=point_coords,point_labels=point_labels)
    print("seg_with_promp:",masks.shape)
    pil_images = draw_bitmask(imput_image,masks)
    return masks,pil_images

def seg_all(imput_image):
    if isinstance(imput_image, Image.Image):
        imput_image = pil_image_to_numpy(imput_image)

    sam = sam_model_registry["vit_b"](checkpoint="./sam_vit_b_01ec64.pth")
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(imput_image)
    pil_images = draw_bitmask(imput_image,masks)
    # pil_images = draw_polygon(imput_image,masks)
    # pil_images = draw_bitmask_split(imput_image,masks)
    return masks,pil_images

# 为每个二值掩码生成一张图片
def draw_bitmask_split(np_image,masks):
    for i,obj in enumerate(masks):
        print("segmentation:",obj["segmentation"].shape)
        view = Visualizer(np_image)
        view.draw_binary_mask(obj["segmentation"])
        vis_image = view.get_output()
        pil_images = visimage_to_pil([vis_image],idx=i)
    return pil_images

# 绘制二值掩码
def draw_bitmask(np_image,masks):
    view = Visualizer(np_image)
    for obj in masks:
        if "segmentation" in obj:
            print("segmentation:",obj["segmentation"].shape)
            view.draw_binary_mask(obj["segmentation"])
        else:
            view.draw_binary_mask(obj)
        
    vis_image = view.get_output()
    pil_images = visimage_to_pil([vis_image])
    return pil_images

# 绘制多边形掩码
def draw_polygon(np_image,masks):
    view = Visualizer(np_image)
    for obj in masks:
        polygon = bitmask_to_polygon(obj["segmentation"])
        view.draw_polygon(polygon,"k")
    vis_image = view.get_output()
    pil_images = visimage_to_pil([vis_image])
    return pil_images


# 二值掩码转换为多边形掩码
def bitmask_to_polygon(mask):
    col_mask = np.asfortranarray(mask)
    contours = measure.find_contours(col_mask,0.5)
    print("contours------",contours.shape)
    for i,contour in enumerate(contours):
        contour = np.flip(contour, axis=1)
        print(f"polygon_{i}",contour.shape)
        # polygon = contour.ravel().tolist()
        # print(f"polygon_{i}",polygon)
    return contour

# VIS图片转换为pil
def visimage_to_pil(visimages,need_save=False,idx=0):
    pil_images = []
    for i,visimage in enumerate(visimages):
        visualized_image = visimage.get_image()[:, :, ::-1]
        pil_image = Image.fromarray(visualized_image)
        if need_save:
            pil_image.save(f"{idx}_{i}.jpg")
        pil_images.append(pil_image)
    return pil_images

def image_to_mask(image, threshold=128):
    # 将图像转换为灰度图像
    if image.mode != 'L':
        image = image.convert('L')
    
    # 将像素值映射到二进制值
    mask_array = np.array(image) > threshold
    
    # 创建一个与原始图像大小相同的数组，用映射后的二进制值填充
    mask_image = Image.fromarray(np.uint8(mask_array) * 255)
    
    return mask_image
class SamAnything:
    _instance = None
    _lock = threading.Lock()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(SamAnything, cls).__new__(cls)
                    cls._instance._initialize(*args, **kwargs)
        return cls._instance

    def _initialize(self, checkpoint_path="./sam_vit_b_01ec64.pth"):
        self.sam = sam_model_registry["vit_b"](checkpoint=checkpoint_path).to(self.device)
        self.predictor = SamPredictor(self.sam)
        self.mask_generator = SamAutomaticMaskGenerator(self.sam)

    def seg_with_promp(self, imput_image, point_coords=None, box=None):
        if isinstance(imput_image, Image.Image):
            imput_image = pil_image_to_numpy(imput_image)
        point_labels = None
        if point_coords is not None:
            point_labels = np.ones(point_coords.shape[0])

        self.predictor.set_image(imput_image)
        masks = None

        if box is not None:
            masks, _, _ = self.predictor.predict(box=box)
        elif point_coords is not None and point_labels is not None:
            masks, _, _ = self.predictor.predict(point_coords=point_coords, point_labels=point_labels)

        print("seg_with_promp:", masks.shape)
        pil_images = self.draw_bitmask(imput_image, masks)
        return masks, pil_images

    def seg_all(self, imput_image):
        if isinstance(imput_image, Image.Image):
            imput_image = pil_image_to_numpy(imput_image)

        masks = self.mask_generator.generate(imput_image)
        pil_images = self.draw_bitmask(imput_image, masks)
        return masks, pil_images

    @staticmethod
    def draw_bitmask_split(np_image, masks):
        pil_images = []
        for i, obj in enumerate(masks):
            print("segmentation:", obj["segmentation"].shape)
            view = Visualizer(np_image)
            view.draw_binary_mask(obj["segmentation"])
            vis_image = view.get_output()
            pil_images.extend(SamAnything.visimage_to_pil([vis_image], idx=i))
        return pil_images

    @staticmethod
    def draw_bitmask(np_image, masks):
        view = Visualizer(np_image)
        for obj in masks:
            if "segmentation" in obj:
                print("segmentation:", obj["segmentation"].shape)
                view.draw_binary_mask(obj["segmentation"])
            else:
                view.draw_binary_mask(obj)
        
        vis_image = view.get_output()
        pil_images = SamAnything.visimage_to_pil([vis_image])
        return pil_images

    @staticmethod
    def draw_polygon(np_image, masks):
        view = Visualizer(np_image)
        for obj in masks:
            polygon = SamAnything.bitmask_to_polygon(obj["segmentation"])
            view.draw_polygon(polygon, "k")
        vis_image = view.get_output()
        pil_images = SamAnything.visimage_to_pil([vis_image])
        return pil_images

    @staticmethod
    def bitmask_to_polygon(mask):
        col_mask = np.asfortranarray(mask)
        contours = measure.find_contours(col_mask, 0.5)
        print("contours------", len(contours))
        for i, contour in enumerate(contours):
            contour = np.flip(contour, axis=1)
            print(f"polygon_{i}", contour.shape)
        return contours

    @staticmethod
    def visimage_to_pil(visimages, need_save=True, idx=0):
        pil_images = []
        for i, visimage in enumerate(visimages):
            visualized_image = visimage.get_image()
            pil_image = Image.fromarray(visualized_image)
            if need_save:
                pil_image.save(f"{idx}_{i}.jpg")
            pil_images.append(pil_image)
        return pil_images

    @staticmethod
    def image_to_mask(image, threshold=128):
        if image.mode != 'L':
            image = image.convert('L')
        mask_array = np.array(image) > threshold
        mask_image = Image.fromarray(np.uint8(mask_array) * 255)
        return mask_image

# if __name__ == "__main__":
#     np_image = read_image("./test/face1.jpeg")
#     print("np_image:",np_image.shape)
#     SamAnything.initialize_sam("./sam_vit_b_01ec64.pth")
#     SamAnything.seg_all(np_image)