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