通过SAM2自助一个分割图像的工具

不知道各位有没有碰到需要对一个图中的物体进行分割操作的需求时，却没有一个趁手的工具。或者你想要一个正版的PS?

小编也碰到过同样的问题，不过好在现在


AI


的强大，就让它按照我的要求写了一个分割图像的程序：通过SAM2(Segment Anython Model-2)，实现把物体从背景中分割，并且可以把分割图像使用的掩码、以及从图像中获取的物体图像处理为透明背景的PNG图后输出到指定的目录中。

先看视频，后查代码。

下面是实际操作的示意。

记住几个简单的操作：

按住鼠标左键拖拉矩形框，选择需要分割的图像区域

单击鼠标左键，添加额外关注的局部区域

单击鼠标右键，减少不需要的局部区域

Save Path设置保存输出图文件的路径

点击Export按钮，输出分割完成的图到Save Path

这里设想各位已经安装好了SAM2，并且也下载了对应的模型权重文件。

请留意其中的三个保存不同文件的路径。

"""
本python脚本包含或使用了 SAM2（Segment Anything Model 2）相关技术，用于实现图像分割、目标识别辅助或视觉处理能力。该技术可能由第三方提供，
并受其相应的开源许可证、版权声明及使用条款约束。 本产品与 SAM2 的原始开发者之间不存在隶属、赞助或官方认证关系，
除非另有明确说明。 使用本功能时，系统输出可能不准确、遗漏或偏差，用户应结合实际场景进行复核.
"""


importos
importjson
importdatetime
importtkinterastk
fromtkinterimportfiledialog
importcv2
importnumpyasnp
importtorch
importmatplotlib.pyplotasplt
frommatplotlib.widgetsimportButton, RectangleSelector
fromPILimportImage
fromsam2.build_samimportbuild_sam2
fromsam2.sam2_image_predictorimportSAM2ImagePredictor


classSAM2InteractiveSelector:
def__init__(self, checkpoint, model_cfg):
"""
SAM2 interactive GUI:
- Left drag on empty area = draw new box
- Left short click on empty area = foreground point
- Right short click = background point
- Click on existing box center = move box
- Click on existing box corner/edge = resize box
- Click inside box = still add FG point
"""
self.checkpoint = checkpoint
self.model_cfg = model_cfg
self.output_dir ="Path to your defaul save or export dir/img"
os.makedirs(self.output_dir, exist_ok=True)
self.device = torch.device("cuda"iftorch.cuda.is_available()else"cpu")
# Model
self.model =None
self.predictor =None
self.model_loaded =False
# Image
self.image_path =None
self.image_bgr =None
self.image_rgb =None
self.original_image =None
self.image_loaded =False
# Segmentation state
self.rect_coords =None
self.current_mask =None
self.current_rgba =None# 存储当前显示的RGBA图像
# Prompt points
self.point_coords_list = []
self.point_labels_list = []
# Mouse gesture state
self.mouse_press_pos =None
self.mouse_press_button =None
self.is_dragging =False
self.drag_threshold =5
# Box interaction state
self.box_selected =False
self.box_action =None # None / move / resize
self.box_handle =None # tl/tr/bl/br/top/bottom/left/right
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
self.load_sam2_model()
self.setup_ui()
plt.show()
# ---------------------- Model ----------------------
defload_sam2_model(self):
print("Loading SAM2 model...")
try:
self.model = build_sam2(self.model_cfg,self.checkpoint, device=self.device)
self.model.eval()
self.predictor = SAM2ImagePredictor(self.model)
self.model_loaded =True
print(f"SAM2 model loaded successfully on{self.device}")
exceptExceptionase:
print(f"Model loading failed:{e}")
raise
# ---------------------- UI ----------------------
defsetup_ui(self):
self.fig = plt.figure(figsize=(16,9))
# Set window title
try:
self.fig.canvas.manager.set_window_title("Amphenol Sensors - SAM2 Interactive Selector")
exceptException:
pass
# 主图像显示区域
self.ax_img = plt.axes([0.05,0.12,0.58,0.83]) #(left, bottom, width, height)
self.ax_img.axis("off")
self.ax_img.set_title("Main View - Click Open Image to load a picture", fontsize=12)
btn_x =0.70
btn_w =0.15
btn_h =0.04
ax_open = plt.axes([btn_x,0.94, btn_w, btn_h])
ax_savepath = plt.axes([btn_x,0.89, btn_w, btn_h])
ax_save_params = plt.axes([btn_x,0.84, btn_w, btn_h])
ax_clear_points = plt.axes([btn_x,0.79, btn_w, btn_h])
ax_confirm = plt.axes([btn_x,0.73, btn_w, btn_h])
ax_cancel = plt.axes([btn_x,0.68, btn_w, btn_h])
ax_export = plt.axes([btn_x,0.63, btn_w, btn_h])
ax_reset = plt.axes([btn_x,0.58, btn_w, btn_h])
ax_quit = plt.axes([btn_x,0.53, btn_w, btn_h])
# 结果预览窗口（带棋格背景）
self.ax_result = plt.axes([btn_x,0.08, btn_w,0.42])
self.ax_result.axis("off")
self.ax_result.set_title("Result Preview (Checkerboard BG)", fontsize=10)
# 初始化棋格背景预览
self.init_checkerboard_preview()
ax_info = plt.axes([btn_x,0.01, btn_w,0.05])
ax_info.axis("off")
self.info_text = ax_info.text(
0.5,0.5,
"Ready
Open an image first",
ha="center", va="center",
fontsize=8, wrap=True
)
self.btn_open = Button(ax_open,"Open Image", color="lightyellow", hovercolor="yellow")
self.btn_savepath = Button(ax_savepath,"Save Path", color="lightyellow", hovercolor="yellow")
self.btn_save_params = Button(ax_save_params,"Save Box+Points", color="lavender", hovercolor="plum")
self.btn_clear_points = Button(ax_clear_points,"Clear All Points", color="mistyrose", hovercolor="lightcoral")
self.btn_confirm = Button(ax_confirm,"Confirm", color="lightgreen", hovercolor="green")
self.btn_cancel = Button(ax_cancel,"Cancel", color="lightcoral", hovercolor="red")
self.btn_export = Button(ax_export,"Export PNG", color="lightblue", hovercolor="blue")
self.btn_reset = Button(ax_reset,"Reset", color="lightgray", hovercolor="gray")
self.btn_quit = Button(ax_quit,"Quit", color="salmon", hovercolor="red")
self.btn_open.on_clicked(self.open_image_clicked)
self.btn_savepath.on_clicked(self.select_output_directory_btn)
self.btn_save_params.on_clicked(self.save_current_params)
self.btn_clear_points.on_clicked(self.clear_all_points)
self.btn_confirm.on_clicked(self.confirm_selection)
self.btn_cancel.on_clicked(self.cancel_selection)
self.btn_export.on_clicked(self.export_png)
self.btn_reset.on_clicked(self.reset_all)
self.btn_quit.on_clicked(self.quit_program)
# 只负责空白区域拖拽创建框，不让它参与框编辑
self.rect_selector = RectangleSelector(
self.ax_img,
self.on_rect_select,
useblit=True,
button=[1],
minspanx=5,
minspany=5,
spancoords="pixels",
interactive=False,
props=dict(facecolor="none", edgecolor="red", alpha=1.0, fill=False)
)
self.fig.canvas.mpl_connect("button_press_event",self.on_mouse_press)
self.fig.canvas.mpl_connect("motion_notify_event",self.on_mouse_move)
self.fig.canvas.mpl_connect("button_release_event",self.on_mouse_release)
self.fig.canvas.mpl_connect("key_press_event",self.on_key_press)


definit_checkerboard_preview(self):
"""初始化棋格背景预览区域"""
self.ax_result.clear()
self.ax_result.axis("off")
self.ax_result.set_title("Result Preview (Checkerboard BG)", fontsize=10)

# 创建一个简单的棋格背景示例
preview_size =100
checkerboard =self.create_checkerboard(preview_size, preview_size,10)
self.ax_result.imshow(checkerboard)
self.ax_result.text(0.5,0.5,"No result yet
Run segmentation first",
ha="center", va="center", transform=self.ax_result.transAxes,
fontsize=10, color="gray")
self.fig.canvas.draw_idle()


defcreate_checkerboard(self, height, width, tile_size=10):
"""创建棋格背景图像"""
checker = np.zeros((height, width,3), dtype=np.uint8)
foryinrange(height):
forxinrange(width):
if((x // tile_size) + (y // tile_size)) %2==0:
checker[y, x] = [220,220,220] # 浅灰色
else:
checker[y, x] = [255,255,255] # 白色
returnchecker


defupdate_result_preview(self, rgba_image):
"""更新结果预览窗口，显示带棋格背景的透明图像"""
ifrgba_imageisNone:
return

try:
self.ax_result.clear()
self.ax_result.axis("off")
self.ax_result.set_title("Result Preview (Checkerboard BG)", fontsize=10)

# 将RGBA图像与棋格背景混合
preview =self.blend_with_checkerboard(rgba_image)

self.ax_result.imshow(preview)
self.fig.canvas.draw_idle()
exceptExceptionase:
print(f"Update result preview failed:{e}")
defblend_with_checkerboard(self, rgba_image, tile_size=20):
"""将RGBA图像与棋格背景混合"""
h, w = rgba_image.shape[:2]

# 创建棋格背景
checkerboard = np.zeros((h, w,3), dtype=np.uint8)
foryinrange(h):
forxinrange(w):
if((x // tile_size) + (y // tile_size)) %2==0:
checkerboard[y, x] = [220,220,220]
else:
checkerboard[y, x] = [255,255,255]

# 提取alpha通道
alpha = rgba_image[...,3:4] /255.0
rgb = rgba_image[..., :3]

# 混合
blended = (rgb * alpha + checkerboard * (1- alpha)).astype(np.uint8)
returnblended


defupdate_info(self, message):
ifhasattr(self,"info_text"):
self.info_text.set_text(message)
self.fig.canvas.draw_idle()
# ---------------------- File dialogs ----------------------
defselect_image_file(self):
root = tk.Tk()
root.withdraw()
file_path = filedialog.askopenfilename(
title="Select image file",
filetypes=[
("Image files","*.jpg *.jpeg *.png *.bmp *.tiff *.tif"),
("All files","*.*")
]
)
root.destroy()
returnfile_path


defselect_output_directory(self):
root = tk.Tk()
root.withdraw()
directory = filedialog.askdirectory(
title="Select save directory",
initialdir=self.output_dir
)
root.destroy()
ifdirectory:
self.output_dir = directory
os.makedirs(self.output_dir, exist_ok=True)
returnTrue
returnFalse
# ---------------------- Image load ----------------------
defopen_image_clicked(self, event):
path =self.select_image_file()
ifnotpath:
self.update_info("No image selected")
return
ifself.load_image(path):
self.update_info("Image loaded.
Drag box and add points.")
else:
self.update_info("Failed to load image")
defselect_output_directory_btn(self, event):
ok =self.select_output_directory()
ifok:
self.update_info(f"Save path:
{self.output_dir}")
else:
self.update_info(f"Using current path:
{self.output_dir}")


defload_image(self, image_path):
try:
self.image_path = image_path
self.image_bgr = cv2.imread(image_path, cv2.IMREAD_COLOR)
ifself.image_bgrisNone:
print(f"Cannot load image:{image_path}")
returnFalse
self.image_rgb = cv2.cvtColor(self.image_bgr, cv2.COLOR_BGR2RGB)
self.original_image =self.image_rgb.copy()
self.predictor.set_image(self.image_rgb)
self.image_loaded =True
self.rect_coords =None
self.current_mask =None
self.current_rgba =None
self.point_coords_list = []
self.point_labels_list = []
self.box_selected =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
print(f"Image loaded successfully! Size:{self.image_rgb.shape}")
self.display_image()
self.init_checkerboard_preview() # 重置预览窗口
returnTrue
exceptExceptionase:
print(f"Image loading failed:{e}")
returnFalse
defdisplay_image(self):
self.ax_img.clear()
self.ax_img.imshow(self.original_image)
self.ax_img.set_title(
"Left click = FG point | Right click = BG point | Left drag = box | Click box edge/center = edit",
fontsize=12
)
self.ax_img.axis("off")
self.fig.canvas.draw_idle()
# ---------------------- Box helpers ----------------------
defget_box_hit_test(self, x, y, box, tol=8):
ifboxisNone:
returnNone
x_min, y_min, x_max, y_max = box
cx = (x_min + x_max) /2.0
cy = (y_min + y_max) /2.0
corners = {
"tl": (x_min, y_min),
"tr": (x_max, y_min),
"bl": (x_min, y_max),
"br": (x_max, y_max),
}
forname, (px, py)incorners.items():
ifabs(x - px) <= tol and abs(y - py) <= tol:
                return ("corner", name)
        edges = {
            "top": ((x_min + x_max) / 2.0, y_min),
            "bottom": ((x_min + x_max) / 2.0, y_max),
            "left": (x_min, (y_min + y_max) / 2.0),
            "right": (x_max, (y_min + y_max) / 2.0),
        }
        for name, (px, py) in edges.items():
            if abs(x - px) <= tol and abs(y - py) <= tol:
                return ("edge", name)
        if abs(x - cx) <= tol and abs(y - cy) <= tol:
            return ("center", None)
        if x_min <= x <= x_max and y_min <= y <= y_max:
            return ("inside", None)
        return None


    def move_box_by_delta(self, dx, dy):
        if self.rect_coords is None or self.box_origin is None:
            return
        x_min, y_min, x_max, y_max = self.box_origin
        h, w = self.image_rgb.shape[:2]
        bw = x_max - x_min
        bh = y_max - y_min
        new_x_min = int(round(x_min + dx))
        new_y_min = int(round(y_min + dy))
        new_x_max = int(round(x_max + dx))
        new_y_max = int(round(y_max + dy))
        if new_x_min < 0:
            new_x_min = 0
            new_x_max = bw
        if new_y_min < 0:
            new_y_min = 0
            new_y_max = bh
        if new_x_max >= w:
new_x_max = w -1
new_x_min = new_x_max - bw
ifnew_y_max >= h:
new_y_max = h -1
new_y_min = new_y_max - bh
self.rect_coords = (
max(0,min(new_x_min, w -1)),
max(0,min(new_y_min, h -1)),
max(0,min(new_x_max, w -1)),
max(0,min(new_y_max, h -1)),
)


defresize_box(self, x, y):
ifself.rect_coordsisNoneorself.box_originisNoneorself.box_handleisNone:
return
x_min, y_min, x_max, y_max =self.box_origin
h, w =self.image_rgb.shape[:2]
min_size =5
ifself.box_handle =="tl":
x_min =min(max(0, x), x_max - min_size)
y_min =min(max(0, y), y_max - min_size)
elifself.box_handle =="tr":
x_max =max(min(w -1, x), x_min + min_size)
y_min =min(max(0, y), y_max - min_size)
elifself.box_handle =="bl":
x_min =min(max(0, x), x_max - min_size)
y_max =max(min(h -1, y), y_min + min_size)
elifself.box_handle =="br":
x_max =max(min(w -1, x), x_min + min_size)
y_max =max(min(h -1, y), y_min + min_size)
elifself.box_handle =="left":
x_min =min(max(0, x), x_max - min_size)
elifself.box_handle =="right":
x_max =max(min(w -1, x), x_min + min_size)
elifself.box_handle =="top":
y_min =min(max(0, y), y_max - min_size)
elifself.box_handle =="bottom":
y_max =max(min(h -1, y), y_min + min_size)
self.rect_coords = (int(x_min),int(y_min),int(x_max),int(y_max))


defdraw_box_and_handles(self):
ifself.rect_coordsisNone:
return
x_min, y_min, x_max, y_max =self.rect_coords
edge_color ="yellow"ifself.box_selectedelse"red"
line_w =3ifself.box_selectedelse2
rect = plt.Rectangle(
(x_min, y_min),
x_max - x_min,
y_max - y_min,
fill=False,
edgecolor=edge_color,
linewidth=line_w
)
self.ax_img.add_patch(rect)
ifself.box_selected:
handles = [
(x_min, y_min), (x_max, y_min), (x_min, y_max), (x_max, y_max),
((x_min + x_max) /2.0, y_min),
((x_min + x_max) /2.0, y_max),
(x_min, (y_min + y_max) /2.0),
(x_max, (y_min + y_max) /2.0),
((x_min + x_max) /2.0, (y_min + y_max) /2.0),
]
forhx, hyinhandles:
self.ax_img.scatter([hx], [hy], s=55, c="white", edgecolors="black", zorder=5)
# ---------------------- Box selection ----------------------
defon_rect_select(self, eclick, erelease):
ifnotself.image_loaded:
return
ifself.box_editing:
return
ifeclick.xdataisNoneoreclick.ydataisNoneorerelease.xdataisNoneorerelease.ydataisNone:
return
try:
x1, y1 =int(eclick.xdata),int(eclick.ydata)
x2, y2 =int(erelease.xdata),int(erelease.ydata)
x_min, x_max =min(x1, x2),max(x1, x2)
y_min, y_max =min(y1, y2),max(y1, y2)
h, w =self.image_rgb.shape[:2]
x_min =max(0,min(x_min, w -1))
x_max =max(0,min(x_max, w -1))
y_min =max(0,min(y_min, h -1))
y_max =max(0,min(y_max, h -1))
if(x_max - x_min) < 5 or (y_max - y_min) < 5:
                print("Box too small, ignored.")
                return
            self.rect_coords = (x_min, y_min, x_max, y_max)
            self.box_selected = True
            self.box_action = None
            self.box_handle = None
            self.box_origin = None
            self.box_press_pos = None
            print(f"Selected box: ({x_min}, {y_min}) -> ({x_max},{y_max})")
self.update_info(f"Box selected:
({x_min},{y_min}) -> ({x_max},{y_max})")
self.redraw_current_state()
exceptExceptionase:
print(f"Error in rectangle selection:{e}")
# ---------------------- Mouse gesture: click vs drag ----------------------
defon_mouse_press(self, event):
ifnotself.image_loaded:
return
ifevent.inaxes !=self.ax_img:
return
ifevent.xdataisNoneorevent.ydataisNone:
return
ifevent.buttonin[1,3]:
self.mouse_press_pos = (event.xdata, event.ydata)
self.mouse_press_button = event.button
self.is_dragging =False
ifevent.button ==1andself.rect_coordsisnotNone:
x, y =int(event.xdata),int(event.ydata)
hit =self.get_box_hit_test(x, y,self.rect_coords, tol=8)
self.box_action =None
self.box_handle =None
self.box_origin =None
self.box_press_pos =None
self.box_editing =False
ifhitisnotNone:
kind, name = hit
self.box_selected =True
# 只有 edge/corner 才 resize，center 才 move
ifkind =="corner":
self.box_action ="resize"
self.box_handle = name
self.box_origin =self.rect_coords
self.box_press_pos = (event.xdata, event.ydata)
self.box_editing =True
self.update_info("Resize box by dragging corner.")
elifkind =="edge":
self.box_action ="resize"
self.box_handle = name
self.box_origin =self.rect_coords
self.box_press_pos = (event.xdata, event.ydata)
self.box_editing =True
self.update_info("Resize box by dragging edge.")
elifkind =="center":
self.box_action ="move"
self.box_origin =self.rect_coords
self.box_press_pos = (event.xdata, event.ydata)
self.box_editing =True
self.update_info("Move box by dragging center.")
else:
# inside：不进入编辑，保持为普通单击加点
self.box_action =None
self.redraw_current_state()


defon_mouse_move(self, event):
ifself.mouse_press_posisNone:
return
ifevent.inaxes !=self.ax_img:
return
ifevent.xdataisNoneorevent.ydataisNone:
return
dx = event.xdata -self.mouse_press_pos[0]
dy = event.ydata -self.mouse_press_pos[1]
dist = (dx * dx + dy * dy) **0.5
ifdist >=self.drag_threshold:
self.is_dragging =True
ifself.box_editingandself.mouse_press_button ==1:
ifself.box_action =="move"andself.box_press_posisnotNone:
move_dx = event.xdata -self.box_press_pos[0]
move_dy = event.ydata -self.box_press_pos[1]
self.move_box_by_delta(move_dx, move_dy)
self.redraw_current_state()
elifself.box_action =="resize":
self.resize_box(int(event.xdata),int(event.ydata))
self.redraw_current_state()


defon_mouse_release(self, event):
ifnotself.image_loaded:
self.mouse_press_pos =None
self.mouse_press_button =None
self.is_dragging =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
return
ifself.mouse_press_posisNone:
return
ifevent.inaxes !=self.ax_imgorevent.xdataisNoneorevent.ydataisNone:
self.mouse_press_pos =None
self.mouse_press_button =None
self.is_dragging =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
return
x, y =int(event.xdata),int(event.ydata)
# 如果当前是在编辑框，则释放时只结束编辑，不再进入点选逻辑
ifself.mouse_press_button ==1andself.box_editing:
self.box_selected =True
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
self.mouse_press_pos =None
self.mouse_press_button =None
self.is_dragging =False
self.update_info("Box updated.")
self.redraw_current_state()
return
# 左键短按：空白或框内都应该先尝试按"选点"
ifself.mouse_press_button ==1:
ifnotself.is_dragging:
ifself.rect_coordsisnotNone:
hit =self.get_box_hit_test(x, y,self.rect_coords, tol=8)
# 只有 corner/edge/center 才是框编辑入口
# inside 仍然当作 FG 点
ifhitisnotNone:
kind, _ = hit
ifkindin("corner","edge","center"):
self.box_selected =True
self.redraw_current_state()
else:
self.add_point(x, y, label=1)
else:
self.add_point(x, y, label=1)
else:
self.add_point(x, y, label=1)
elifself.mouse_press_button ==3:
ifnotself.is_dragging:
self.add_point(x, y, label=0)
self.mouse_press_pos =None
self.mouse_press_button =None
self.is_dragging =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
defadd_point(self, x, y, label):
self.point_coords_list.append([x, y])
self.point_labels_list.append(label)
print(f"Added point: ({x},{y}), label={label}")
self.redraw_current_state()
color ="green"iflabel ==1else"red"
marker ="*"iflabel ==1else"x"
label_name ="FG"iflabel ==1else"BG"
self.ax_img.scatter(
[x], [y],
color=color,
marker=marker,
s=200iflabel ==1else120,
edgecolor="white"iflabel ==1elseNone,
linewidth=1.0iflabel ==1else2.0
)
self.fig.canvas.draw_idle()
print(f"Added{label_name}point: ({x},{y})")
self.update_info(f"Added{label_name}point:
({x},{y})")
# ---------------------- Save / clear params ----------------------
defsave_current_params(self, event):
ifnotself.image_loaded:
self.update_info("Load an image first")
return
params = {
"image_path":self.image_path,
"timestamp": datetime.datetime.now().isoformat(),
"box":self.rect_coords,
"points": [
{"x":int(p[0]),"y":int(p[1]),"label":int(l)}
forp, linzip(self.point_coords_list,self.point_labels_list)
]
}
try:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
json_path = os.path.join(self.output_dir,f"box_points_{timestamp}.json")
withopen(json_path,"w", encoding="utf-8")asf:
json.dump(params, f, ensure_ascii=False, indent=2)
print(f"Saved parameters to:{json_path}")
self.update_info(f"Saved params:
box_points_{timestamp}.json")
exceptExceptionase:
print(f"Save params failed:{e}")
self.update_info(f"Save params failed:
{e}")
defclear_all_points(self, event):
ifnotself.image_loaded:
self.update_info("Load an image first")
return
self.point_coords_list = []
self.point_labels_list = []
self.redraw_current_state()
self.update_info("All points cleared")
print("All points cleared")
# ---------------------- Keyboard ----------------------
defon_key_press(self, event):
ifevent.key =="enter":
self.confirm_selection(None)
elifevent.key =="escape":
self.cancel_selection(None)
elifevent.key =="c":
self.clear_all_points(None)
# ---------------------- SAM2 prompts ----------------------
defget_prompt_inputs(self):
ifself.rect_coordsisNone:
returnNone,None,None
x_min, y_min, x_max, y_max =self.rect_coords
input_box = np.array([x_min, y_min, x_max, y_max], dtype=np.float32)
point_coords =None
point_labels =None
iflen(self.point_coords_list) >0:
point_coords = np.array(self.point_coords_list, dtype=np.float32)
point_labels = np.array(self.point_labels_list, dtype=np.int32)
returninput_box, point_coords, point_labels
defconfirm_selection(self, event):
ifnotself.image_loaded:
self.update_info("Load an image first")
return
ifself.rect_coordsisNone:
self.update_info("Draw a box first")
return
input_box, point_coords, point_labels =self.get_prompt_inputs()
try:
self.update_info("Segmenting...")
print("Performing SAM2 segmentation with box + points...")
withtorch.inference_mode():
masks, scores, logits =self.predictor.predict(
point_coords=point_coords,
point_labels=point_labels,
box=input_box[None, :],
multimask_output=True
)
best_idx =int(np.argmax(scores))
self.current_mask = masks[best_idx].astype(bool)
# 创建RGBA图像并更新预览
self.current_rgba = np.zeros((self.image_rgb.shape[0],self.image_rgb.shape[1],4), dtype=np.uint8)
self.current_rgba[..., :3] =self.image_rgb
self.current_rgba[...,3] =self.current_mask.astype(np.uint8) *255

# 更新结果预览窗口
self.update_result_preview(self.current_rgba)
self.visualize_result(scores[best_idx])
print(f"Segmentation complete! Score:{scores[best_idx]:.3f}")
self.update_info(f"Done!
Score:{scores[best_idx]:.3f}")
exceptExceptionase:
print(f"Segmentation failed:{e}")
self.update_info(f"Failed:
{e}")
# ---------------------- Display ----------------------
defredraw_current_state(self):
self.ax_img.clear()
self.ax_img.imshow(self.original_image)
ifself.current_maskisnotNone:
overlay = np.zeros((*self.current_mask.shape,4), dtype=np.uint8)
overlay[...,0] =30
overlay[...,1] =144
overlay[...,2] =255
overlay[...,3] =self.current_mask.astype(np.uint8) *120
self.ax_img.imshow(overlay)
self.draw_box_and_handles()
iflen(self.point_coords_list) >0:
pts = np.array(self.point_coords_list)
labels = np.array(self.point_labels_list)
fg = pts[labels ==1]
bg = pts[labels ==0]
iflen(fg) >0:
self.ax_img.scatter(
fg[:,0], fg[:,1],
color="green", marker="*",
s=200, edgecolor="white", linewidth=1.0
)
iflen(bg) >0:
self.ax_img.scatter(
bg[:,0], bg[:,1],
color="red", marker="x",
s=120, linewidth=2.0
)
ifself.rect_coordsisnotNone:
x_min, y_min, x_max, y_max =self.rect_coords
ifself.box_selected:
self.ax_img.set_title(
f"Box selected: ({x_min},{y_min}) -> ({x_max},{y_max})",
fontsize=12
)
else:
self.ax_img.set_title(
"Left click = FG point | Right click = BG point | Left drag = box",
fontsize=12
)
self.ax_img.axis("off")
self.fig.canvas.draw_idle()


defvisualize_result(self, score=None):
self.redraw_current_state()
# ---------------------- Reset / cancel ----------------------
defcancel_selection(self, event):
self.current_mask =None
self.current_rgba =None
self.rect_coords =None
self.point_coords_list = []
self.point_labels_list = []
self.box_selected =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
ifself.image_loaded:
self.display_image()
self.init_checkerboard_preview() # 重置预览窗口
self.update_info("Selection cancelled")


defreset_all(self, event):
self.current_mask =None
self.current_rgba =None
self.rect_coords =None
self.point_coords_list = []
self.point_labels_list = []
self.box_selected =False
self.box_action =None
self.box_handle =None
self.box_press_pos =None
self.box_origin =None
self.box_editing =False
ifself.image_loaded:
self.display_image()
self.init_checkerboard_preview() # 重置预览窗口
self.update_info("Reset done")
# ---------------------- Export ----------------------
defexport_png(self, event):
ifnotself.image_loaded:
self.update_info("Load an image first")
return
ifself.current_maskisNone:
self.update_info("No mask to export")
return
try:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
mask =self.current_mask.astype(np.uint8)
rgba = np.zeros((self.image_rgb.shape[0],self.image_rgb.shape[1],4), dtype=np.uint8)
rgba[..., :3] =self.image_rgb
rgba[...,3] = mask *255
out_name =f"segmented_object_{timestamp}.png"
out_path = os.path.join(self.output_dir, out_name)
Image.fromarray(rgba, mode="RGBA").save(out_path)
mask_path = os.path.join(self.output_dir,f"mask_{timestamp}.png")
cv2.imwrite(mask_path, mask *255)
print(f"Saved transparent PNG:{out_path}")
print(f"Saved binary mask:{mask_path}")
self.update_info(f"Saved:
{out_name}")
self.ax_img.set_title(f"Exported:{out_name}", fontsize=11, color="green")
self.fig.canvas.draw_idle()
exceptExceptionase:
print(f"Export failed:{e}")
self.update_info(f"Export failed:
{e}")


defquit_program(self, event):
print("Closing program...")
plt.close("all")


defmain():
checkpoint ="Path to your dir/sam2.1_hiera_base_plus.pt"
model_cfg ="Path to your dir/sam2.1/sam2.1_hiera_b+.yaml"
try:
SAM2InteractiveSelector(checkpoint, model_cfg)
exceptExceptionase:
print(f"Program error:{e}")
importtraceback
traceback.print_exc()


if__name__ =="__main__":
main()