SAM/segment.py

60 lines
1.4 KiB
Python

import cv2
import os
import numpy as np
from segment_anything import sam_model_registry, SamPredictor
input_dir = 'scripts/input/images'
output_dir = 'scripts/output/mask'
crop_mode = True
print('最好是每加一个点就按w键predict一次')
os.makedirs(output_dir, exist_ok=True)
image_files = [f for f in os.listdir(input_dir) if
f.lower().endswith(('.png', '.jpg', '.jpeg', '.JPG', '.JPEG', '.PNG', '.tiff'))]
sam = sam_model_registry["vit_b"](checkpoint=r"D:\Program Files\Pycharm items\segment-anything-model\weights\vit_b.pth")
_ = sam.to(device="cuda")
predictor = SamPredictor(sam)
WINDOW_WIDTH = 1280
WINDOW_HEIGHT = 720
cv2.namedWindow("image", cv2.WINDOW_NORMAL)
cv2.resizeWindow("image", WINDOW_WIDTH, WINDOW_HEIGHT)
cv2.moveWindow("image", (1920 - WINDOW_WIDTH) // 2, (1080 - WINDOW_HEIGHT) // 2)
def button_image_open():
# 实现打开图片功能
pass
def button_image_init():
# 实现重置选择功能
pass
def button_image_shang():
# 实现切换至上一张图像功能
pass
def button_image_xia():
# 实现切换至下一张图像功能
pass
def button_image_exit():
# 实现退出程序功能
pass
def button_image_Transparency():
# 实现调整透明度功能
pass
def button_image_copymask():
# 实现复制掩码功能
pass
def button_image_saveimg():
# 实现保存图片功能
pass
def button_image_Fusionimg():
# 实现融合背景图片功能
pass