SAM/salt/GUI.py

import sys
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import QTimer
import cv2
import os
import numpy as np
from segment_anything import sam_model_registry, SamPredictor

input_dir = r'D:\Program Files\Pycharm items\segment-anything-model\scripts\input\images'
output_dir = r'D:\Program Files\Pycharm items\segment-anything-model\scripts\output\maskt'
crop_mode = True

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

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1170, 486)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.pushButton_w = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_w.setGeometry(QtCore.QRect(10, 90, 151, 51))
        self.pushButton_w.setObjectName("pushButton_w")
        self.pushButton_a = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_a.setGeometry(QtCore.QRect(10, 160, 151, 51))
        self.pushButton_a.setObjectName("pushButton_a")
        self.pushButton_d = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_d.setGeometry(QtCore.QRect(10, 230, 151, 51))
        self.pushButton_d.setObjectName("pushButton_d")
        self.pushButton_s = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_s.setGeometry(QtCore.QRect(10, 300, 151, 51))
        self.pushButton_s.setObjectName("pushButton_s")
        self.pushButton_q = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_q.setGeometry(QtCore.QRect(10, 370, 151, 51))
        self.pushButton_q.setObjectName("pushButton_q")
        self.label_orign = QtWidgets.QLabel(self.centralwidget)
        self.label_orign.setGeometry(QtCore.QRect(180, 20, 471, 401))
        self.label_orign.setStyleSheet("background-color: rgb(255, 255, 255);")
        self.label_orign.setObjectName("label_orign")
        self.label_pre = QtWidgets.QLabel(self.centralwidget)
        self.label_pre.setGeometry(QtCore.QRect(660, 20, 471, 401))
        self.label_pre.setStyleSheet("background-color: rgb(255, 255, 255);")
        self.label_pre.setObjectName("label_pre")
        self.pushButton_opimg = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_opimg.setGeometry(QtCore.QRect(10, 20, 151, 51))
        self.pushButton_opimg.setObjectName("pushButton_opimg")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 1170, 26))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

        self.pushButton_opimg.clicked.connect(self.open_image)
        self.pushButton_w.clicked.connect(self.predict_image)

        self.timer = QTimer()
        self.timer.timeout.connect(self.update_original_image)
        self.timer.start(100)  # Update every 100 milliseconds

        self.image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.JPG', '.JPEG', '.PNG', '.tiff'))]
        self.current_index = 0
        self.input_point = []
        self.input_label = []
        self.input_stop = False

        cv2.namedWindow("image")
        cv2.setMouseCallback("image", self.mouse_click)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.pushButton_w.setText(_translate("MainWindow", "w"))
        self.pushButton_a.setText(_translate("MainWindow", "a"))
        self.pushButton_d.setText(_translate("MainWindow", "d"))
        self.pushButton_s.setText(_translate("MainWindow", "s"))
        self.pushButton_q.setText(_translate("MainWindow", "q"))
        self.label_orign.setText(_translate("MainWindow", "<html><head/><body><p align=\"center\">原始图像</p></body></html>"))
        self.label_pre.setText(_translate("MainWindow", "<html><head/><body><p align=\"center\">预测图像</p></body></html>"))
        self.pushButton_opimg.setText(_translate("MainWindow", "打开图像"))

    def open_image(self):
        filename, _ = QtWidgets.QFileDialog.getOpenFileName(None, "Open Image File", "", "Image files (*.jpg *.png)")
        if filename:
            self.current_index = 0
            pixmap = QPixmap(filename)
            pixmap = pixmap.scaled(self.label_orign.size(), QtCore.Qt.KeepAspectRatio)
            self.label_orign.setPixmap(pixmap)
            self.label_orign.setAlignment(QtCore.Qt.AlignCenter)
            self.input_point = []
            self.input_label = []

    def predict_image(self):
        if self.current_index < len(self.image_files):
            filename = self.image_files[self.current_index]
            image_orign = cv2.imread(os.path.join(input_dir, filename))
            image_display = cv2.cvtColor(image_orign.copy(), cv2.COLOR_BGR2RGB)

            for point, label in zip(self.input_point, self.input_label):
                color = (0, 255, 0) if label == 1 else (0, 0, 255)
                cv2.circle(image_display, tuple(point), 5, color, -1)

            cv2.imshow("image", image_display)

    def update_original_image(self):
        if self.current_index < len(self.image_files):
            filename = self.image_files[self.current_index]
            image_orign = cv2.imread(os.path.join(input_dir, filename))
            image_display = cv2.cvtColor(image_orign.copy(), cv2.COLOR_BGR2RGB)

            for point, label in zip(self.input_point, self.input_label):
                color = (0, 255, 0) if label == 1 else (0, 0, 255)
                cv2.circle(image_display, tuple(point), 5, color, -1)

            height, width, channel = image_display.shape
            bytesPerLine = 3 * width
            qImg = QImage(image_display.data, width, height, bytesPerLine, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(qImg)
            pixmap = pixmap.scaled(self.label_orign.size(), QtCore.Qt.KeepAspectRatio)
            self.label_orign.setPixmap(pixmap)
            self.label_orign.setAlignment(QtCore.Qt.AlignCenter)

    def mouse_click(self, event, x, y, flags, param):
        if not self.input_stop:
            if event == cv2.EVENT_LBUTTONDOWN:
                self.input_point.append([x, y])
                self.input_label.append(1)
            elif event == cv2.EVENT_RBUTTONDOWN:
                self.input_point.append([x, y])
                self.input_label.append(0)
        else:
            if event == cv2.EVENT_LBUTTONDOWN or event == cv2.EVENT_RBUTTONDOWN:
                print('此时不能添加点,按w退出mask选择模式')

def apply_mask(image, mask, alpha_channel=True):
    if alpha_channel:
        alpha = np.zeros_like(image[..., 0])
        alpha[mask == 1] = 255
        image = cv2.merge((image[..., 0], image[..., 1], image[..., 2], alpha))
    else:
        image = np.where(mask[..., None] == 1, image, 0)
    return image


def apply_color_mask(image, mask, color, color_dark=0.5):
    for c in range(3):
        image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - color_dark) + color_dark * color[c], image[:, :, c])
    return image


def get_next_filename(base_path, filename):
    name, ext = os.path.splitext(filename)
    for i in range(1, 101):
        new_name = f"{name}_{i}{ext}"
        if not os.path.exists(os.path.join(base_path, new_name)):
            return new_name
    return None


def save_masked_image(image, mask, output_dir, filename, crop_mode_):
    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)
    filename = filename[:filename.rfind('.')] + '.png'
    new_filename = get_next_filename(output_dir, filename)

    if new_filename:
        if masked_image.shape[-1] == 4:
            cv2.imwrite(os.path.join(output_dir, new_filename), masked_image, [cv2.IMWRITE_PNG_COMPRESSION, 9])
        else:
            cv2.imwrite(os.path.join(output_dir, new_filename), masked_image)
        print(f"Saved as {new_filename}")
    else:
        print("Could not save the image. Too many variations exist.")


if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    MainWindow.show()
    sys.exit(app.exec_())