抠图党福音：教你一键分割图像

import os
import torch
import os.path as osp
import moxing as mox
path = osp.join(os.getcwd(),'segment-anything')
if not os.path.exists(path):
 mox.file.copy_parallel('obs://modelarts-labs-bj4-v2/case_zoo/segment-anything', path)
 if os.path.exists(path):
 print('Download success')
 else:
        raise Exception('Download Failed')
else:
 print("Model Package already exists!") 

%cd segment-anything
!pip install --upgrade pip
!pip install torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1
!pip install opencv-python matplotlib
!python setup.py install
import numpy as np
import matplotlib.pyplot as plt
import cv2
import copy
import torch
import torchvision
print("PyTorch version:", torch.__version__)
print("Torchvision version:", torchvision.__version__)
print("CUDA is available:", torch.cuda.is_available())

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cuda"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
mask_generator = SamAutomaticMaskGenerator(
    model=sam,
    #points_per_side=32,
    #pred_iou_thresh=0.86,
    #stability_score_thresh=0.92,
    #crop_n_layers=1,
    #crop_n_points_downscale_factor=2,
    #min_mask_region_area=100,  # Requires open-cv to run post-processing
)

def show_anns(anns,image):
 segment_image = copy.copy(image)
 segment_image.astype("uint8")
 if len(anns) == 0:
 return
 sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
 for ann in sorted_anns:
        mask_2d = ann['segmentation']
 h,w = mask_2d.shape
        mask_3d_color = np.zeros((h,w,3), dtype=np.uint8)
        mask = (mask_2d!=0).astype(bool)
 rgb = np.random.randint(0, 255, (1, 3), dtype=np.uint8)
        mask_3d_color[mask_2d[:, :] == 1] = rgb
 segment_image[mask] = segment_image[mask] * 0.5 + mask_3d_color[mask] * 0.5
 return segment_image
image = cv2.imread('images/dog.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
masks = mask_generator.generate(image)
segment_image = show_anns(masks,image)
fig = plt.figure(figsize=(25, 10))
ax1 = fig.add_subplot(1, 2, 1)
plt.title('Original image', fontsize=16)
ax1.axis('off')
ax1.imshow(image)
ax2 = fig.add_subplot(1, 2, 2)
plt.title('Segment image', fontsize=16)
ax2.axis('off')
ax2.imshow(segment_image)
plt.show()

def apply_mask(image, mask, alpha_channel=True):#应用并且响应mask
 if alpha_channel:
        alpha = np.zeros_like(image[..., 0])#制作掩体
        alpha[mask == 1] = 255#兴趣地方标记为1，且为白色
        image = cv2.merge((image[..., 0], image[..., 1], image[..., 2], alpha))#融合图像
 else:
        image = np.where(mask[..., None] == 1, image, 0)
 return image
def mask_image(image, mask, crop_mode_=True):#保存掩盖部分的图像（感兴趣的图像）
 if crop_mode_:
        y, x = np.where(mask)
 y_min, y_max, x_min, x_max = y.min(), y.max(), x.min(), x.max()
 cropped_mask = mask[y_min:y_max+1, x_min:x_max+1]
 cropped_image = image[y_min:y_max+1, x_min:x_max+1]
 masked_image = apply_mask(cropped_image, cropped_mask)
 else:
 masked_image = apply_mask(image, mask)
 return masked_image
def save_masked_image(image, filepath):
 if image.shape[-1] == 4:
        cv2.imwrite(filepath, image, [cv2.IMWRITE_PNG_COMPRESSION, 9])
 else:
        cv2.imwrite(filepath, image)
 print(f"Saved as {filepath}")
def save_anns(anns,image,path):
 if len(anns) == 0:
 return
 sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    index = 1
 for ann in sorted_anns:
        mask_2d = ann['segmentation']
 segment_image = copy.copy(image)
 masked_image = mask_image(segment_image, mask_2d)
        filename = str(index) + '.png'
 filepath = os.path.join(path, filename)
 save_masked_image(masked_image, filepath)
        index = index + 1
save_path = 'result/'
if not os.path.exists(save_path):
 os.mkdir(save_path)
image = cv2.imread('images/dog.jpg')
masks = mask_generator.generate(image)
save_anns(masks,image,save_path)

!pip install gradio==3.24.1
def segment_image(image):
    masks = mask_generator.generate(image)
 return show_anns(masks,image)
def show_image(image):
    masks = mask_generator.generate(image)
 if len(masks) == 0:
 return
 sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True)
    index = 1
 image_list = []
 for ann in sorted_anns:
        mask_2d = ann['segmentation']
 segment_image = copy.copy(image)
 masked_image = mask_image(segment_image, mask_2d)
 image_list.append(masked_image)
 return image_list
import gradio as gr
with gr.Blocks() as demo:
 with gr.Row():
 with gr.Column():
 img_in = gr.Image(source='upload')
 with gr.Row():
 segment_button = gr.Button("segment",variant="primary")
 save_button = gr.Button("segment_images",variant="primary")
 with gr.Row():
 with gr.Column():
 img_out = gr.Image()
 with gr.Row():
 result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=6, height='auto')
 segment_button.click(segment_image,
                 inputs= [img_in], 
                 outputs=[img_out])
 save_button.click(show_image,
                 inputs= [img_in], 
                 outputs=[result_gallery])
demo.launch(share=True)